lakehouse_engine.io.writers.jdbc_writer

 1"""Module that defines the behaviour to write to JDBC targets."""
 2
 3from typing import Callable, OrderedDict
 4
 5from pyspark.sql import DataFrame
 6
 7from lakehouse_engine.core.definitions import OutputSpec
 8from lakehouse_engine.io.writer import Writer
 9
10
11class JDBCWriter(Writer):
12    """Class to write to JDBC targets."""
13
14    def __init__(self, output_spec: OutputSpec, df: DataFrame, data: OrderedDict):
15        """Construct JDBCWriter instances.
16
17        Args:
18            output_spec: output specification.
19            df: dataframe to be writen.
20            data: list of all dfs generated on previous steps before writer.
21        """
22        super().__init__(output_spec, df, data)
23
24    def write(self) -> None:
25        """Write data into JDBC target."""
26        if not self._df.isStreaming:
27            self._write_to_jdbc_in_batch_mode(self._df, self._output_spec)
28        else:
29            stream_df = (
30                self._df.writeStream.trigger(
31                    **Writer.get_streaming_trigger(self._output_spec)
32                )
33                .options(
34                    **self._output_spec.options if self._output_spec.options else {}
35                )
36                .foreachBatch(
37                    self._write_transformed_micro_batch(self._output_spec, self._data)
38                )
39                .start()
40            )
41
42            if self._output_spec.streaming_await_termination:
43                stream_df.awaitTermination(
44                    self._output_spec.streaming_await_termination_timeout
45                )
46
47    @staticmethod
48    def _write_to_jdbc_in_batch_mode(df: DataFrame, output_spec: OutputSpec) -> None:
49        """Write to jdbc in batch mode.
50
51        Args:
52            df: dataframe to write.
53            output_spec: output specification.
54        """
55        df.write.format(output_spec.data_format).partitionBy(
56            output_spec.partitions
57        ).options(**output_spec.options if output_spec.options else {}).mode(
58            output_spec.write_type
59        ).save(
60            output_spec.location
61        )
62
63    @staticmethod
64    def _write_transformed_micro_batch(  # type: ignore
65        output_spec: OutputSpec, data: OrderedDict
66    ) -> Callable:
67        """Define how to write a streaming micro batch after transforming it.
68
69        Args:
70            output_spec: output specification.
71            data: list of all dfs generated on previous steps before writer.
72
73        Returns:
74            A function to be executed in the foreachBatch spark write method.
75        """
76
77        def inner(batch_df: DataFrame, batch_id: int) -> None:
78            transformed_df = Writer.get_transformed_micro_batch(
79                output_spec, batch_df, batch_id, data
80            )
81
82            if output_spec.streaming_micro_batch_dq_processors:
83                transformed_df = Writer.run_micro_batch_dq_process(
84                    transformed_df, output_spec.streaming_micro_batch_dq_processors
85                )
86
87            JDBCWriter._write_to_jdbc_in_batch_mode(transformed_df, output_spec)
88
89        return inner