duckdb_pipeline.py

"""DuckDB implementation for `Pipeline` object."""

import logging
from typing import Optional

from duckdb import DuckDBPyConnection, DuckDBPyRelation

import dve.parser.file_handling as fh
from dve.core_engine.backends.base.reference_data import BaseRefDataLoader, ReferenceConfig
from dve.core_engine.backends.implementations.duckdb.auditing import DDBAuditingManager
from dve.core_engine.backends.implementations.duckdb.contract import DuckDBDataContract
from dve.core_engine.backends.implementations.duckdb.duckdb_helpers import duckdb_get_entity_count
from dve.core_engine.backends.implementations.duckdb.reference_data import DuckDBRefDataLoader
from dve.core_engine.backends.implementations.duckdb.rules import DuckDBStepImplementations
from dve.core_engine.models import SubmissionInfo
from dve.core_engine.type_hints import URI
from dve.pipeline.pipeline import BaseDVEPipeline


# pylint: disable=abstract-method
@duckdb_get_entity_count
class DDBDVEPipeline(BaseDVEPipeline):
    """
    Modified Pipeline class for running a DVE Pipeline with Spark
    """

    # pylint: disable=R0913
    def __init__(
        self,
        processed_files_path: URI,
        audit_tables: DDBAuditingManager,
        connection: DuckDBPyConnection,
        rules_path: Optional[URI],
        submitted_files_path: Optional[URI],
        job_run_id: Optional[int] = None,
        logger: Optional[logging.Logger] = None,
    ):
        self._connection = connection
        super().__init__(
            processed_files_path,
            audit_tables,
            DuckDBDataContract(connection=self._connection),
            DuckDBStepImplementations.register_udfs(connection=self._connection),
            rules_path,
            submitted_files_path,
            job_run_id,
            logger,
        )

    def get_reference_data_loader(
        self, reference_data_config: dict[str, ReferenceConfig], **kwargs
    ) -> DuckDBRefDataLoader:
        return DuckDBRefDataLoader(
            connection=self._connection,
            reference_data_config=reference_data_config,
            dataset_config_uri=fh.get_parent(self._rules_path), # type: ignore
            **kwargs
        )

    # pylint: disable=arguments-differ
    def write_file_to_parquet(  # type: ignore
        self, submission_file_uri: URI, submission_info: SubmissionInfo, output: URI
    ):
        return super().write_file_to_parquet(
            submission_file_uri, submission_info, output, DuckDBPyRelation
        )