lakehouse_engine.terminators.cdf_processor

  1"""Defines change data feed processor behaviour."""
  2
  3from datetime import datetime, timedelta
  4from typing import OrderedDict
  5
  6from delta.tables import DeltaTable
  7from pyspark.sql import DataFrame
  8from pyspark.sql.functions import col, date_format
  9
 10from lakehouse_engine.core.definitions import (
 11    InputSpec,
 12    OutputFormat,
 13    OutputSpec,
 14    ReadType,
 15    TerminatorSpec,
 16    WriteType,
 17)
 18from lakehouse_engine.core.exec_env import ExecEnv
 19from lakehouse_engine.io.reader_factory import ReaderFactory
 20from lakehouse_engine.io.writer_factory import WriterFactory
 21from lakehouse_engine.utils.logging_handler import LoggingHandler
 22
 23
 24class CDFProcessor(object):
 25    """Change data feed processor class."""
 26
 27    _logger = LoggingHandler(__name__).get_logger()
 28
 29    @classmethod
 30    def expose_cdf(cls, spec: TerminatorSpec) -> None:
 31        """Expose CDF to external location.
 32
 33        Args:
 34            spec: terminator specification.
 35        """
 36        cls._logger.info("Reading CDF from input table...")
 37
 38        df_cdf = ReaderFactory.get_data(cls._get_table_cdf_input_specs(spec))
 39        new_df_cdf = df_cdf.withColumn(
 40            "_commit_timestamp",
 41            date_format(col("_commit_timestamp"), "yyyyMMddHHmmss"),
 42        )
 43
 44        cls._logger.info("Writing CDF to external table...")
 45        cls._write_cdf_to_external(
 46            spec,
 47            new_df_cdf.repartition(
 48                spec.args.get(
 49                    "materialized_cdf_num_partitions", col("_commit_timestamp")
 50                )
 51            ),
 52        )
 53
 54        # used to delete old data on CDF table (don't remove parquet).
 55        if spec.args.get("clean_cdf", True):
 56            cls._logger.info("Cleaning CDF table...")
 57            cls.delete_old_data(spec)
 58
 59        # used to delete old parquet files.
 60        if spec.args.get("vacuum_cdf", False):
 61            cls._logger.info("Vacuuming CDF table...")
 62            cls.vacuum_cdf_data(spec)
 63
 64    @staticmethod
 65    def _write_cdf_to_external(
 66        spec: TerminatorSpec, df: DataFrame, data: OrderedDict = None
 67    ) -> None:
 68        """Write cdf results dataframe.
 69
 70        Args:
 71            spec: terminator specification.
 72            df: dataframe with cdf results to write.
 73            data: list of all dfs generated on previous steps before writer.
 74        """
 75        WriterFactory.get_writer(
 76            spec=OutputSpec(
 77                spec_id="materialized_cdf",
 78                input_id="input_table",
 79                location=spec.args["materialized_cdf_location"],
 80                write_type=WriteType.APPEND.value,
 81                data_format=spec.args.get("data_format", OutputFormat.DELTAFILES.value),
 82                options=spec.args["materialized_cdf_options"],
 83                partitions=["_commit_timestamp"],
 84            ),
 85            df=df,
 86            data=data,
 87        ).write()
 88
 89    @staticmethod
 90    def _get_table_cdf_input_specs(spec: TerminatorSpec) -> InputSpec:
 91        """Get the input specifications from a terminator spec.
 92
 93        Args:
 94            spec: terminator specifications.
 95
 96        Returns:
 97            List of input specifications.
 98        """
 99        options = {
100            "readChangeFeed": "true",
101            **spec.args.get("db_table_options", {}),
102        }
103
104        input_specs = InputSpec(
105            spec_id="input_table",
106            db_table=spec.args["db_table"],
107            read_type=ReadType.STREAMING.value,
108            data_format=OutputFormat.DELTAFILES.value,
109            options=options,
110        )
111
112        return input_specs
113
114    @classmethod
115    def delete_old_data(cls, spec: TerminatorSpec) -> None:
116        """Delete old data from cdf delta table.
117
118        Args:
119            spec: terminator specifications.
120        """
121        today_datetime = datetime.today()
122        limit_date = today_datetime + timedelta(
123            days=spec.args.get("days_to_keep", 30) * -1
124        )
125        limit_timestamp = limit_date.strftime("%Y%m%d%H%M%S")
126
127        cdf_delta_table = DeltaTable.forPath(
128            ExecEnv.SESSION, spec.args["materialized_cdf_location"]
129        )
130
131        cdf_delta_table.delete(col("_commit_timestamp") < limit_timestamp)
132
133    @classmethod
134    def vacuum_cdf_data(cls, spec: TerminatorSpec) -> None:
135        """Vacuum old data from cdf delta table.
136
137        Args:
138            spec: terminator specifications.
139        """
140        cdf_delta_table = DeltaTable.forPath(
141            ExecEnv.SESSION, spec.args["materialized_cdf_location"]
142        )
143
144        cdf_delta_table.vacuum(spec.args.get("vacuum_hours", 168))