Add PipelineConfig, PipelineFactory, DagFactory

Author: Tomás Link

src/gfw/common/beam/pipeline/dag/factory.py

@@ -0,0 +1,120 @@
+"""Factories for building Apache Beam DAGs with BigQuery integration.
+
+This module defines abstract base classes for DAG factories that produce
+Apache Beam pipelines, including support for creating BigQuery read/write
+clients and helpers with optional mocking capabilities.
+
+Classes:
+    DagFactory: Abstract base class providing BigQuery client factories and
+        requiring a build_dag method.
+
+    LinearDagFactory: Extends DagFactory for linear pipelines composed of
+        sources, core, optional side inputs, and sinks.
+"""
+
+from abc import ABC, abstractmethod
+from functools import partial
+from typing import Callable, Optional, Tuple
+
+from apache_beam import PTransform
+
+from gfw.common.beam.pipeline.dag import Dag, LinearDag
+from gfw.common.beam.transforms import WriteToPartitionedBigQuery
+from gfw.common.bigquery.helper import BigQueryHelper
+from gfw.common.pipeline.config import PipelineConfig
+
+
+class DagFactory(ABC):
+    """Abstract base class for DAG factories producing Apache Beam pipelines.
+
+    Provides factory properties for BigQuery read/write clients and helpers.
+    """
+
+    def __init__(self, config: PipelineConfig) -> None:
+        self.config = config
+
+    @property
+    def write_to_bigquery_factory(self) -> Callable[..., WriteToPartitionedBigQuery]:
+        """Returns a factory for WriteToPartitionedBigQuery clients.
+
+        Uses mocked clients if configured.
+        """
+        return WriteToPartitionedBigQuery.get_client_factory(mocked=self.config.mock_bq_clients)
+
+    @property
+    def bigquery_helper_factory(self) -> Callable[..., BigQueryHelper]:
+        """Returns a factory for BigQueryHelper instances.
+
+        Returns:
+            Callable that creates BigQueryHelper instances with
+            the appropriate client factory.
+        """
+        client_factory = BigQueryHelper.get_client_factory(mocked=self.config.mock_bq_clients)
+        return partial(BigQueryHelper, client_factory=client_factory)
+
+    @abstractmethod
+    def build_dag(self) -> Dag:
+        """Builds the DAG.
+
+        Must be implemented in subclasses.
+
+        Returns:
+            A tuple of PTransforms representing the DAG components.
+        """
+        pass
+
+
+class LinearDagFactory(DagFactory, ABC):
+    """Base class for linear DAG factories that assemble sources, core, side inputs, and sinks."""
+
+    @property
+    @abstractmethod
+    def sources(self) -> Tuple[PTransform, ...]:
+        """Returns the source PTransforms for the LinearDag.
+
+        Returns:
+            Tuple of PTransforms serving as data sources.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def core(self) -> PTransform:
+        """Returns the core PTransform that processes data in the LinearDag.
+
+        Returns:
+            The core processing PTransform.
+        """
+        pass
+
+    @property
+    def side_inputs(self) -> Optional[PTransform]:
+        """Returns optional side inputs PTransform for the LinearDag.
+
+        Returns:
+            A PTransform for side inputs or None if not used.
+        """
+        return None
+
+    @property
+    @abstractmethod
+    def sinks(self) -> Tuple[PTransform, ...]:
+        """Returns the sink PTransforms for the LinearDag.
+
+        Returns:
+            Tuple of PTransforms serving as data sinks.
+        """
+        pass
+
+    def build_dag(self) -> LinearDag:
+        """Builds a LinearDag instance from the configured pipeline parts.
+
+        Returns:
+            A LinearDag composed of sources, core, side inputs, and sinks.
+        """
+        return LinearDag(
+            sources=tuple(self.sources),
+            core=self.core,
+            side_inputs=self.side_inputs,
+            sinks=tuple(self.sinks),
+        )

src/gfw/common/beam/pipeline/factory.py

@@ -0,0 +1,60 @@
+"""Factory for constructing Beam pipelines from configuration and DAG factories.
+
+This module defines the PipelineFactory class, which builds a fully configured
+Pipeline instance from a given PipelineConfig and DagFactory.
+"""
+
+from gfw.common.beam.pipeline import Pipeline
+from gfw.common.beam.pipeline.dag.factory import DagFactory
+from gfw.common.pipeline.config import PipelineConfig
+
+
+class PipelineFactory:
+    """Builds a Beam Pipeline from a configuration object and a DAG factory.
+
+    Attributes:
+        config:
+            Configuration for the pipeline, including version and CLI arguments.
+
+        dag_factory:
+            Factory that produces the pipeline's DAG.
+
+        name:
+            Optional name for the pipeline.
+    """
+
+    def __init__(
+        self,
+        config: PipelineConfig,
+        dag_factory: DagFactory,
+        name: str = "",
+    ) -> None:
+        """Initializes the factory with config, DAG factory, and optional name.
+
+        Args:
+            config:
+                The pipeline configuration.
+
+            dag_factory:
+                Factory that provides the pipeline DAG.
+
+            name:
+                Optional name for the pipeline.
+        """
+        self.config = config
+        self.dag_factory = dag_factory
+        self.name = name
+
+    def build_pipeline(self) -> Pipeline:
+        """Constructs and returns a fully configured Pipeline instance.
+
+        Returns:
+            A pipeline with DAG, version, name, and CLI arguments.
+        """
+        return Pipeline(
+            name=self.name,
+            version=self.config.version,
+            dag=self.dag_factory.build_dag(),
+            unparsed_args=self.config.unknown_unparsed_args,
+            **self.config.unknown_parsed_args,
+        )

src/gfw/common/bigquery/helper.py

@@ -3,7 +3,7 @@
 import logging
 
 from dataclasses import dataclass
-from functools import cached_property, partial
+from functools import cached_property
 from pathlib import Path
 from typing import Any, Callable, Dict, Iterator, List, Optional, Union
 from unittest import mock
@@ -118,8 +118,16 @@ def mocked(cls, **kwargs: Any) -> "BigQueryHelper":
     @classmethod
     def get_client_factory(cls, mocked: bool = False) -> Callable[..., bigquery.client.Client]:
         """Returns a factory for bigquery.Client objects."""
+
+        def mock_client_factory(*args: Any, **kwargs: Any) -> mock.Mock:
+            # Extract project from kwargs or use default None.
+            # Otherwise project is not set and is needed.
+            project = kwargs.get("project", None)
+            client = mock.create_autospec(bigquery.Client, project=project, instance=True)
+            return client
+
         if mocked:
-            return partial(mock.create_autospec, bigquery.client.Client, instance=True)
+            return mock_client_factory
 
         return bigquery.client.Client
 

src/gfw/common/pipeline/__init__.py

@@ -0,0 +1,14 @@
+"""Configuration models and utilities for building data pipelines.
+
+The `pipeline` package defines configuration models and utilities
+for building and managing data processing pipelines.
+
+Modules:
+    config: Contains classes and helpers for parsing and managing
+            pipeline configuration, including date ranges and
+            unknown arguments from the CLI.
+
+This package is typically used by pipeline factories and runners
+to extract configuration parameters, validate inputs, and
+initialize pipeline components with consistent settings.
+"""

src/gfw/common/pipeline/config.py

@@ -0,0 +1,110 @@
+"""This module defines the `PipelineConfig` class used to configure data pipeline executions.
+
+It includes:
+- A dataclass `PipelineConfig` that stores date ranges and any unknown arguments.
+- Automatic parsing of date strings into `datetime.date` objects.
+- Utility methods to convert from a `SimpleNamespace` or to a dictionary.
+- A custom exception `PipelineConfigError` for handling invalid configuration inputs.
+
+Intended for use in CLI-based or programmatic pipeline setups where date ranges
+and additional arguments need to be passed and validated.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import asdict, dataclass, field
+from datetime import date
+from functools import cached_property
+from types import SimpleNamespace
+from typing import Any
+
+
+ERROR_DATE = "Dates must be in ISO format. Got: {}."
+
+
+class PipelineConfigError(Exception):
+    """Custom exception for pipeline configuration errors."""
+
+    pass
+
+
+@dataclass
+class PipelineConfig(ABC):
+    """Configuration object for data pipeline execution.
+
+    Args:
+        date_range:
+            Tuple of start and end dates in ISO format (YYYY-MM-DD).
+
+        mock_bq_clients:
+            If True, all BigQuery interactions will be mocked.
+
+        unknown_parsed_args:
+            Parsed CLI or config arguments not explicitly defined in the config.
+
+        unknown_unparsed_args:
+            Raw unparsed CLI arguments.
+    """
+
+    date_range: tuple[str, str]
+    mock_bq_clients: bool = False
+    unknown_parsed_args: dict[str, Any] = field(default_factory=dict)
+    unknown_unparsed_args: tuple[str, ...] = ()
+
+    @classmethod
+    def from_namespace(cls, ns: SimpleNamespace) -> "PipelineConfig":
+        """Creates a PipelineConfig instance from a SimpleNamespace.
+
+        Args:
+            ns: Namespace containing attributes matching PipelineConfig fields.
+
+        Returns:
+            A new PipelineConfig instance.
+        """
+        return cls(**vars(ns))
+
+    @cached_property
+    def parsed_date_range(self) -> tuple[date, date]:
+        """Returns the parsed start and end dates as `datetime.date` objects.
+
+        Raises:
+            PipelineConfigError: If any of the dates are not in valid ISO format.
+
+        Returns:
+            A tuple containing start and end dates as `date` objects.
+        """
+        try:
+            start_str, end_str = self.date_range
+            return (date.fromisoformat(start_str), date.fromisoformat(end_str))
+        except ValueError as e:
+            raise PipelineConfigError(ERROR_DATE.format(self.date_range)) from e
+
+    @property
+    def start_date(self) -> date:
+        """Returns the start date of the configured range.
+
+        Returns:
+            A `date` object representing the start of the range.
+        """
+        return self.parsed_date_range[0]
+
+    @property
+    def end_date(self) -> date:
+        """Returns the end date of the configured range.
+
+        Returns:
+            A `date` object representing the end of the range.
+        """
+        return self.parsed_date_range[1]
+
+    @abstractmethod
+    @property
+    def version(self) -> str:
+        """Returns the version of the pipeline."""
+
+    def to_dict(self) -> dict[str, Any]:
+        """Converts the PipelineConfig to a dictionary.
+
+        Returns:
+            A dictionary representation of the configuration.
+        """
+        return asdict(self)