src/common/materializer/dependent_dags.py

# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Defines utility functions to help construct task dependent DAGs."""

from collections import defaultdict
import dataclasses
import datetime
import graphlib
import logging
from pathlib import Path
import textwrap
from typing import Optional

from common.materializer import dag_types
from common.materializer import generate_assets
from common.py_libs import constants
from common.py_libs import cortex_bq_client
from common.py_libs import dag_generator
from common.py_libs import jinja


def _get_parent_task_ids(parent_files: list[str]) -> list[str]:
    task_ids = []
    for p in parent_files:
        if p == "start_task":
            task_ids.append(p)
        else:
            task_ids.append(f"refresh_{Path(p).stem}")
    return task_ids


def _get_dag_setting(obj: dag_types.BqObject) -> Optional[dag_types.Dag]:
    if not (obj.table_setting and obj.table_setting.dag_setting):
        return None
    return obj.table_setting.dag_setting


def get_validated_load_freq(
        top_level_objs: dict[str, dag_types.BqObject]) -> str:
    """Return validated load_frequency defined by top level nodes.

    Top level nodes must have a consistent schedule defined. Multiple nodes
    can define the same load_frequence, or nodes can leave it unset as long as
    on of the top level nodes has it set.

    Args:
        top_level_objs: map of top level BQ objects that are direct children of
            the DAG root keyed by sql file name.

    Returns: Validated load_frequency string defining DAG schedule.

    Raises:
        ValueError: if the top level nodes don't have any defined load_frequency
            or if they conflict.
    """
    schedule = None
    for sql_file, obj in top_level_objs.items():
        assert obj.table_setting is not None

        if obj.table_setting.load_frequency is None:
            continue

        load_frequency = obj.table_setting.load_frequency
        if schedule is None:
            schedule = load_frequency
            continue

        if load_frequency != schedule:
            raise ValueError(
                "Top level node deteced with conflicting load_frequency. "
                f"{sql_file} set to '{load_frequency}' but expected "
                f"'{schedule}'. Please resolve any load_frequency conflicts "
                f"among top level nodes: {list(top_level_objs.keys())}")

    if schedule is None:
        raise ValueError(
            "At least one top level node must define load_frequency. Top level "
            f"nodes include: {list(top_level_objs.keys())}")
    return schedule


def get_task_deps(bq_object_settings: dict) -> dict[str, dag_types.BqObject]:
    """Get a dictionary of task dependent objects from reporting settings.

    Any object that has dag_setting defined is considered task dependent, even
    if there are no defined edges.

    Args:
        bq_object_settings: A dict representing BQ [in]dependent objects
            defined in a reportings settings yaml file.

    Returns:
        A dict of task dependent BQ objects keyed by their SQL file.
    """
    reporting_objects = dag_types.ReportingObjects.from_dict(bq_object_settings)

    # Populate map of BQ objects and set of objects that are parent or children
    # deps.
    obj_map = {}
    dep_files = set()
    all_objs = []
    if reporting_objects.bq_independent_objects:
        all_objs.extend(reporting_objects.bq_independent_objects)
    if reporting_objects.bq_dependent_objects:
        all_objs.extend(reporting_objects.bq_dependent_objects)

    for obj in all_objs:
        obj_map[obj.sql_file] = obj

        # Objects without dag_setting are ignored and considered independent.
        dag_setting = _get_dag_setting(obj)
        if not dag_setting:
            continue

        dep_files.add(obj.sql_file)
        if dag_setting.parents:
            dep_files = dep_files.union(dag_setting.parents)

    return {k: v for k, v in obj_map.items() if k in dep_files}


@dataclasses.dataclass
class DependentDagGenerator:
    """Generates DAG files with task dependencies.

    Attributes:
        module_name: Name of module (e.g. "cm360", "sap")
        target_dataset_name: Name of BQ dataset - e.g. "my_project.my_dataset"
        target_dataset_type: Type of dataset - e.g. "reporting" or "cdc".
        allow_telemetry: Cortex config file option to enable telemetry.
        location: Location used for BigQueryInsertJob operators.
        output_dir: Directory to write the generated files.
        jinja_data_file: File containing jinja substitutions used to create
            generated SQL.
    """
    module_name: str
    target_dataset_name: str
    target_dataset_type: str
    allow_telemetry: bool
    location: str
    output_dir: Path
    jinja_data_file: Path

    def _generate_sql_file(self, sql_file: str, dag_name: str) -> Path:
        """Returns the relative path to the generated table refresh sql file."""

        table_name = Path(sql_file).stem
        full_table_name = f"{self.target_dataset_name}.{table_name}"
        bq_client = cortex_bq_client.CortexBQClient()

        # Generate core sql text from sql file after applying Jinja parameters.
        core_sql = jinja.apply_jinja_params_to_file(sql_file,
                                                    str(self.jinja_data_file))
        generate_assets.validate_sql(bq_client, core_sql)
        refresh_query = generate_assets.generate_table_refresh_sql(
            bq_client, full_table_name, core_sql)

        relative_sql_file = Path(table_name).with_suffix(".sql")
        full_sql_file = self.output_dir / dag_name / relative_sql_file
        full_sql_file.parent.mkdir(exist_ok=True, parents=True)
        full_sql_file.write_text(refresh_query, encoding="utf-8")
        logging.info("Generated DAG SQL file : %s", full_sql_file)

        return relative_sql_file

    def _generate_dag_file(self, dag_name: str, ordered_nodes: list[str],
                           task_dep_objs: dict[str, dag_types.BqObject],
                           load_frequency: str) -> Path:
        logging.info("Generating DAG file for: %s", dag_name)

        parent_dir = Path(__file__).resolve().parent
        template_dir = parent_dir / "templates"
        header_template = (template_dir /
                           "airflow_task_dep_dag_template_reporting.py")
        bq_op_template = template_dir / "bq_insert_job_template.txt"
        dag_full_name = "_".join(
            [self.target_dataset_name.replace(".", "_"), dag_name])
        today = datetime.datetime.now()

        # General template substitutions.
        subs = {
            "dag_full_name": dag_full_name,
            "module_name": self.module_name,
            "tgt_dataset_type": self.target_dataset_type,
            "load_frequency": load_frequency,
            "year": today.year,
            "month": today.month,
            "day": today.day,
            "runtime_labels_dict": "",  # A place holder for label dict string,
            "bq_location": self.location
        }
        if self.allow_telemetry:
            subs["runtime_labels_dict"] = str(constants.CORTEX_JOB_LABEL)
        if self.target_dataset_type == "reporting":
            subs["tags"] = [self.module_name, self.target_dataset_type]

        # Create DAG header.
        dag_header = dag_generator.generate_str_from_template(
            header_template, **subs)

        # Create BQ ops and edges.
        bq_op_strs = []
        edge_strs = []
        for sql_file in ordered_nodes:
            # Start tasks don't have parent edges or BQ operators.
            if sql_file == "start_task":
                continue

            table_name = Path(sql_file).stem
            if sql_file == "stop_task":
                task_id = sql_file
            else:
                task_id = f"refresh_{table_name}"

            # Generate edges.
            obj = task_dep_objs[sql_file]
            parents = _get_parent_task_ids(
                obj.table_setting.dag_setting.parents)  # type: ignore
            assert parents is not None
            if len(parents) > 1:
                parents_str = ", ".join(sorted(parents))
                edge_str = f"[{parents_str}] >> {task_id}"
            else:
                edge_str = f"{parents[0]} >> {task_id}"

            edge_str = textwrap.indent(edge_str, " " * 4)
            edge_strs.append(edge_str)

            # Don't generate BQ operator for stop tasks .
            if sql_file == "stop_task":
                continue

            generated_sql_file = self._generate_sql_file(sql_file, dag_name)

            # SQL File specific substitutions.
            subs["table_name"] = table_name
            subs["query_file"] = generated_sql_file

            # Generate BQ operators.
            bq_op = dag_generator.generate_str_from_template(
                bq_op_template, **subs)
            bq_op = textwrap.indent(bq_op, " " * 4)
            bq_op_strs.append(bq_op)

        generated_dag = "\n\n".join([dag_header, *bq_op_strs, *edge_strs])
        generated_dag_file = (self.output_dir / dag_name /
                              dag_full_name).with_suffix(".py")
        generated_dag_file.parent.mkdir(parents=True, exist_ok=True)
        generated_dag_file.write_text(generated_dag)

        logging.info("Generated DAG py file : %s", generated_dag_file)
        return generated_dag_file

    def create_dep_dag(self, dag_name: str,
                       task_dep_objs: dict[str, dag_types.BqObject]) -> Path:
        """Creates a single DAG from a map of task dependent objects.

        Args:
            dag_name: Name of the DAG to generate.
            task_dep_objs: map of sql file names to their BQ Object settings.
                This argument is mutated to add a "start_task" parent to all top
                level nodes and a "stop_task" key with leaf nodes as parents.

        Returns: Path to generated DAG file.

        Raises:
            ValueError: If a cycle is detected in the DAG.
            RuntimeError: If a task dependent object is provided without
                dag_setting defined.
        """
        topo_sorter = graphlib.TopologicalSorter()

        # Top level nodes only have the root as a parent dependency
        top_level_nodes = {}
        leaf_nodes = set(task_dep_objs.keys())

        for sql_file, obj in task_dep_objs.items():
            dag_setting = _get_dag_setting(obj)
            if not dag_setting:
                raise RuntimeError(
                    "Task dependent object must have dag_setting defined: "
                    f"{obj}")

            # Add an inferred root parent for top level nodes.
            if not dag_setting.parents:
                top_level_nodes[sql_file] = obj
                dag_setting.parents = ["start_task"]

            leaf_nodes = leaf_nodes.difference(dag_setting.parents)
            topo_sorter.add(sql_file, *dag_setting.parents)

        # Add an inferred stop node
        topo_sorter.add("stop_task", *leaf_nodes)
        task_dep_objs["stop_task"] = dag_types.BqObject(
            type=dag_types.BqObjectType.BQ_OBJECT_TYPE_UNSPECIFIED,
            table_setting=dag_types.Table(dag_setting=dag_types.Dag(
                name=dag_name, parents=list(leaf_nodes))))

        load_freq = get_validated_load_freq(top_level_nodes)

        try:
            ordered_nodes = [*topo_sorter.static_order()]
        except graphlib.CycleError as e:
            raise ValueError(f"Cyclic dependency detected in DAG: {dag_name}. "
                             "DAG must not have cycles.") from e

        return self._generate_dag_file(dag_name, ordered_nodes, task_dep_objs,
                                       load_freq)

    def create_dep_dags(
            self, task_dep_objs: dict[str, dag_types.BqObject]) -> list[Path]:
        """Creates a series of DAGs from a dict of objs keyed by sql file.

        Returns: List of paths to generated DAG files.
        """
        dags = defaultdict(dict)
        output_files = []

        # Group objects into DAGs by name.
        for sql_file, obj in task_dep_objs.items():
            if not (obj.table_setting and obj.table_setting.dag_setting and
                    obj.table_setting.dag_setting.name):
                raise RuntimeError("Task dependent object was detected but "
                                   f"dag_setting.name is undefined: {obj}")

            dag_name = obj.table_setting.dag_setting.name
            dag_name = dag_name.lower().replace(" ", "_")
            dags[dag_name][sql_file] = obj

        for dag_name, obj in dags.items():
            output_files.append(self.create_dep_dag(dag_name, obj))

        return output_files