diff --git a/CHANGELOG.md b/CHANGELOG.md index 0173a9e..d43038f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,10 @@ cp ../shared/helpers.py payload/helpers.py ``` +### Added + +- **`StreamingClient.write_dmo_deltas(name, dataframe)`** – start a streaming query that writes each micro-batch to the target DMO and return the `StreamingQuery` handle. It is the DMO counterpart to `write_dlo_deltas` and pairs with `read_dmo_deltas()`. + ## 6.1.0 ### Added diff --git a/README.md b/README.md index e7df12a..aa979fa 100644 --- a/README.md +++ b/README.md @@ -158,6 +158,7 @@ For a streaming (delta) transform, use `StreamingClient`, which exposes the stre * `read_dlo_deltas()` – Read the streaming change feed (deltas) of a Data Lake Object as a streaming DataFrame. * `read_dmo_deltas()` – Read the streaming change feed (deltas) of a Data Model Object as a streaming DataFrame. * `write_dlo_deltas(name, spark_dataframe)` – Write a streaming DataFrame of deltas to a Data Lake Object; returns the started `StreamingQuery` +* `write_dmo_deltas(name, spark_dataframe)` – Write a streaming DataFrame of deltas to a Data Model Object; returns the started `StreamingQuery` `find_file_path`, `llm_gateway_generate_text`, and `einstein_predict` are available on both clients. @@ -222,6 +223,15 @@ query = client.write_dlo_deltas("Output__dll", transformed) query.awaitTermination() ``` +A streaming transform can also read and write Data Model Objects with `read_dmo_deltas()` and `write_dmo_deltas()`: + +```python +deltas = client.read_dmo_deltas() +transformed = deltas.withColumn("description__c", upper(col("description__c"))) +query = client.write_dmo_deltas("Output__dlm", transformed) +query.awaitTermination() +``` + Notes: - These methods only run inside the Data Cloud streaming (`DELTA_SYNC`) runtime. Locally (`datacustomcode run`) they raise `NotImplementedError`, since there is no change feed to stream. diff --git a/src/datacustomcode/client.py b/src/datacustomcode/client.py index 49d1407..69f7ea1 100644 --- a/src/datacustomcode/client.py +++ b/src/datacustomcode/client.py @@ -726,6 +726,25 @@ def write_dlo_deltas( self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO) return self._writer.write_dlo_deltas(name, dataframe, **kwargs) # type: ignore[no-any-return] + def write_dmo_deltas( + self, name: str, dataframe: PySparkDataFrame, **kwargs + ) -> StreamingQuery: + """Write a streaming DataFrame of deltas to a DMO in Data Cloud. + + Starts a streaming query that writes each micro-batch to the + target DMO and returns the ``StreamingQuery`` handle; the caller + typically calls ``query.awaitTermination()``. + + Args: + name: The name of the DMO to write to. + dataframe: The streaming PySpark DataFrame to write. + + Returns: + The started ``StreamingQuery``. + """ + self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DLO) + return self._writer.write_dmo_deltas(name, dataframe, **kwargs) # type: ignore[no-any-return] + def auto_write_to_dlo(self, name: str, dataframe: PySparkDataFrame) -> None: """Write a PySpark DataFrame to a DLO in Data Cloud automatically picking the WriteMode. diff --git a/src/datacustomcode/deploy.py b/src/datacustomcode/deploy.py index 2354bdf..51d6cdb 100644 --- a/src/datacustomcode/deploy.py +++ b/src/datacustomcode/deploy.py @@ -415,13 +415,7 @@ def is_streaming(self) -> bool: def _validate_layers(self) -> "DataTransformConfig": read_is_dlo = isinstance(self.permissions.read, DloPermission) write_is_dlo = isinstance(self.permissions.write, DloPermission) - if self.is_streaming: - if not write_is_dlo: - raise ValueError( - "A streaming transform must write to a DLO " - "(permissions.write must be a 'dlo' entry)." - ) - elif read_is_dlo != write_is_dlo: + if read_is_dlo != write_is_dlo: raise ValueError( "permissions.read and permissions.write must both reference " "DLOs or both reference DMOs (got " diff --git a/src/datacustomcode/io/writer/base.py b/src/datacustomcode/io/writer/base.py index d24a33b..56b89f9 100644 --- a/src/datacustomcode/io/writer/base.py +++ b/src/datacustomcode/io/writer/base.py @@ -87,7 +87,7 @@ def write_dlo_deltas( Args: name: Target Data Lake Object name. dataframe: Streaming PySpark DataFrame produced from a - ``read_dlo_deltas`` / ``read_dmo_deltas`` source. + ``read_dlo_deltas`` source. Returns: The started ``StreamingQuery``; the caller drives its lifecycle @@ -102,3 +102,35 @@ def write_dlo_deltas( "streaming runtime; the local writer does not support streaming " "deltas." ) + + def write_dmo_deltas( + self, name: str, dataframe: PySparkDataFrame + ) -> StreamingQuery: + """Write a streaming DataFrame of deltas to a Data Model Object. + + Streaming counterpart to :meth:`write_to_dmo`. Starts a streaming query + that writes each micro-batch to the target DMO via the Data Cloud + streaming sink and returns the resulting ``StreamingQuery`` handle. The + runtime owns the trigger and checkpoint location; callers pass only the + table name. Concrete streaming behavior is provided by the deployed + Data Cloud runtime; the base implementation raises + :class:`NotImplementedError`. + + Args: + name: Target Data Model Object name. + dataframe: Streaming PySpark DataFrame produced from a + ``read_dmo_deltas`` source. + + Returns: + The started ``StreamingQuery``; the caller drives its lifecycle + (typically ``query.awaitTermination()``). + + Raises: + NotImplementedError: If the active writer does not support streaming + deltas (e.g. the local development writers). + """ + raise NotImplementedError( + "write_dmo_deltas is only supported when running in the Data Cloud " + "streaming runtime; the local writer does not support streaming " + "deltas." + ) diff --git a/src/datacustomcode/scan.py b/src/datacustomcode/scan.py index 059925a..a3bf718 100644 --- a/src/datacustomcode/scan.py +++ b/src/datacustomcode/scan.py @@ -34,7 +34,7 @@ DATA_ACCESS_METHODS = ["read_dlo", "read_dmo", "write_to_dlo", "write_to_dmo"] STREAMING_READ_METHODS = ["read_dlo_deltas", "read_dmo_deltas"] -STREAMING_WRITE_METHODS = ["write_dlo_deltas"] +STREAMING_WRITE_METHODS = ["write_dlo_deltas", "write_dmo_deltas"] DATA_TRANSFORM_CONFIG_TEMPLATE = { "sdkVersion": get_version(), @@ -181,6 +181,7 @@ class StreamingDataAccessLayerCalls(pydantic.BaseModel): read_dlo_deltas: bool read_dmo_deltas: bool write_dlo_deltas: frozenset[str] + write_dmo_deltas: frozenset[str] @pydantic.model_validator(mode="after") def validate_access_layer(self) -> StreamingDataAccessLayerCalls: @@ -193,10 +194,19 @@ def validate_access_layer(self) -> StreamingDataAccessLayerCalls: "A streaming transform must read from at least one DLO or DMO " "delta stream (read_dlo_deltas / read_dmo_deltas)." ) - if not self.write_dlo_deltas: + if not self.write_dlo_deltas and not self.write_dmo_deltas: raise ValueError( - "A streaming transform must write to at least one DLO via " - "write_dlo_deltas." + "A streaming transform must write to at least one DLO or DMO " + "via write_dlo_deltas / write_dmo_deltas." + ) + if self.write_dlo_deltas and self.write_dmo_deltas: + raise ValueError( + "Cannot write DLO and DMO deltas in the same streaming transform." + ) + if self.read_layer != self.write_layer: + raise ValueError( + "A streaming transform must read and write the same layer " + f"(got read={self.read_layer}, write={self.write_layer})." ) return self @@ -205,6 +215,11 @@ def read_layer(self) -> str: """Return the read source layer, ``"dlo"`` or ``"dmo"``.""" return "dlo" if self.read_dlo_deltas else "dmo" + @property + def write_layer(self) -> str: + """Return the write target layer, ``"dlo"`` or ``"dmo"``.""" + return "dlo" if self.write_dlo_deltas else "dmo" + class ClientMethodVisitor(ast.NodeVisitor): """AST Visitor that finds all instances of Client read/write method calls.""" @@ -217,6 +232,7 @@ def __init__(self) -> None: self._read_dlo_deltas: bool = False self._read_dmo_deltas: bool = False self._write_dlo_deltas_instances: set[str] = set() + self._write_dmo_deltas_instances: set[str] = set() self.variable_values: Dict[str, Union[str, None]] = {} def visit_Assign(self, node: ast.Assign) -> None: @@ -261,6 +277,8 @@ def visit_Call(self, node: ast.Call) -> None: name = self._resolve_name_arg(node.args[0]) if name and method_name == "write_dlo_deltas": self._write_dlo_deltas_instances.add(name) + elif name and method_name == "write_dmo_deltas": + self._write_dmo_deltas_instances.add(name) self.generic_visit(node) def _resolve_name_arg(self, arg: ast.expr) -> Union[str, None]: @@ -278,6 +296,7 @@ def is_streaming(self) -> bool: self._read_dlo_deltas or self._read_dmo_deltas or bool(self._write_dlo_deltas_instances) + or bool(self._write_dmo_deltas_instances) ) def found(self) -> DataAccessLayerCalls: @@ -293,6 +312,7 @@ def found_streaming(self) -> StreamingDataAccessLayerCalls: read_dlo_deltas=self._read_dlo_deltas, read_dmo_deltas=self._read_dmo_deltas, write_dlo_deltas=frozenset(self._write_dlo_deltas_instances), + write_dmo_deltas=frozenset(self._write_dmo_deltas_instances), ) @@ -516,10 +536,10 @@ def _update_streaming_config(existing_config: dict[str, Any], file_path: str) -> ) read_names = [source_name] if source_name else [] - write_names = list(output.write_dlo_deltas) + write_names = list(output.write_dlo_deltas or output.write_dmo_deltas) existing_config["permissions"] = { "read": {read_layer: read_names}, - "write": {"dlo": write_names}, + "write": {output.write_layer: write_names}, } diff --git a/tests/io/writer/test_print.py b/tests/io/writer/test_print.py index a10f2a3..258a765 100644 --- a/tests/io/writer/test_print.py +++ b/tests/io/writer/test_print.py @@ -65,6 +65,12 @@ def test_write_dlo_deltas_not_supported_locally(self, print_writer, mock_datafra print_writer.write_dlo_deltas("test_dll", mock_dataframe) assert "write_dlo_deltas" in str(exc_info.value) + def test_write_dmo_deltas_not_supported_locally(self, print_writer, mock_dataframe): + """Streaming delta writes are not supported by the local writer.""" + with pytest.raises(NotImplementedError) as exc_info: + print_writer.write_dmo_deltas("test_dlm", mock_dataframe) + assert "write_dmo_deltas" in str(exc_info.value) + def test_config_name(self): """Test that the CONFIG_NAME class variable is set correctly.""" assert PrintDataCloudWriter.CONFIG_NAME == "PrintDataCloudWriter" diff --git a/tests/test_client.py b/tests/test_client.py index 1293ea1..748c0fc 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -440,6 +440,40 @@ def test_write_dlo_deltas_after_dmo_read_raises_exception( assert "test_dmo" in str(exc_info.value) writer.write_dlo_deltas.assert_not_called() + def test_write_dmo_deltas(self, reset_client, mock_spark): + reader = MagicMock(spec=BaseDataCloudReader) + writer = MagicMock(spec=BaseDataCloudWriter) + mock_df = MagicMock(spec=DataFrame) + mock_query = MagicMock() + writer.write_dmo_deltas.return_value = mock_query + + client = StreamingClient(reader=reader, writer=writer) + client._record_dmo_access("some_dmo") + + result = client.write_dmo_deltas("test_dmo", mock_df, extra_param=True) + + writer.write_dmo_deltas.assert_called_once_with( + "test_dmo", mock_df, extra_param=True + ) + assert result is mock_query + + def test_write_dmo_deltas_after_dlo_read_raises_exception( + self, reset_client, mock_spark + ): + """Streaming DMO write is subject to the same DLO/DMO mixing guard.""" + reader = MagicMock(spec=BaseDataCloudReader) + writer = MagicMock(spec=BaseDataCloudWriter) + mock_df = MagicMock(spec=DataFrame) + + client = StreamingClient(reader=reader, writer=writer) + client._record_dlo_access("test_dlo") + + with pytest.raises(DataCloudAccessLayerException) as exc_info: + client.write_dmo_deltas("test_dmo", mock_df) + + assert "test_dlo" in str(exc_info.value) + writer.write_dmo_deltas.assert_not_called() + def test_streaming_read_write_flow(self, reset_client, mock_spark): """A read_dlo_deltas → write_dlo_deltas flow stays within the DLO layer.""" reader = MagicMock(spec=BaseDataCloudReader) @@ -458,6 +492,24 @@ def test_streaming_read_write_flow(self, reset_client, mock_spark): writer.write_dlo_deltas.assert_called_once_with("target_dll", stream_df) assert "source_dll" in client._data_layer_history[DataCloudObjectType.DLO] + def test_streaming_dmo_read_write_flow(self, reset_client, mock_spark): + """A read_dmo_deltas → write_dmo_deltas flow stays within the DMO layer.""" + reader = MagicMock(spec=BaseDataCloudReader) + writer = MagicMock(spec=BaseDataCloudWriter) + stream_df = MagicMock(spec=DataFrame) + reader.read_dmo_deltas.return_value = stream_df + + client = StreamingClient(reader=reader, writer=writer) + + with patch("datacustomcode.client.config") as mock_config: + mock_config.streaming_source = "source_dlm" + df = client.read_dmo_deltas() + client.write_dmo_deltas("target_dlm", df) + + reader.read_dmo_deltas.assert_called_once_with() + writer.write_dmo_deltas.assert_called_once_with("target_dlm", stream_df) + assert "source_dlm" in client._data_layer_history[DataCloudObjectType.DMO] + def test_auto_write_to_dlo(self, reset_client, mock_spark): reader = MagicMock(spec=BaseDataCloudReader) writer = MagicMock(spec=BaseDataCloudWriter) diff --git a/tests/test_deploy.py b/tests/test_deploy.py index 49bcc46..dc167a0 100644 --- a/tests/test_deploy.py +++ b/tests/test_deploy.py @@ -1683,9 +1683,7 @@ def test_is_streaming_false_when_source_absent(self): ) assert config.is_streaming is False - def test_streaming_allows_dmo_read_dlo_write(self): - """Streaming may read a DMO change feed and write a DLO (writes are - DLO-only)""" + def test_streaming_allows_dmo_read_dmo_write(self): config = DataTransformConfig( sdkVersion="1.0.0", entryPoint="entrypoint.py", @@ -1693,24 +1691,24 @@ def test_streaming_allows_dmo_read_dlo_write(self): streamingSource=StreamingSource(type="dmo", name="input_dmo__dlm"), permissions=Permissions( read=DmoPermission(dmo=["input_dmo__dlm"]), - write=DloPermission(dlo=["output_dlo"]), + write=DmoPermission(dmo=["output_dmo__dlm"]), ), ) assert config.is_streaming is True assert isinstance(config.permissions.read, DmoPermission) - assert isinstance(config.permissions.write, DloPermission) + assert isinstance(config.permissions.write, DmoPermission) - def test_streaming_rejects_dmo_write(self): - """Streaming writes must target a DLO.""" - with pytest.raises(ValueError, match="must write to a DLO"): + def test_streaming_rejects_dmo_read_dlo_write(self): + """Streaming follows the same same-layer rule as batch.""" + with pytest.raises(ValueError, match="both reference"): DataTransformConfig( sdkVersion="1.0.0", entryPoint="entrypoint.py", dataspace="default", - streamingSource=StreamingSource(type="dlo", name="input_dlo"), + streamingSource=StreamingSource(type="dmo", name="input_dmo__dlm"), permissions=Permissions( - read=DloPermission(dlo=["input_dlo"]), - write=DmoPermission(dmo=["output_dmo__dlm"]), + read=DmoPermission(dmo=["input_dmo__dlm"]), + write=DloPermission(dlo=["output_dlo"]), ), ) diff --git a/tests/test_scan.py b/tests/test_scan.py index 716c233..71ead13 100644 --- a/tests/test_scan.py +++ b/tests/test_scan.py @@ -735,7 +735,7 @@ def my_function(event, context): client = StreamingClient() deltas = client.read_dmo_deltas() - query = client.write_dlo_deltas("Account_copy__dll", deltas) + query = client.write_dmo_deltas("Account_copy__dlm", deltas) query.awaitTermination() """ ) @@ -751,6 +751,13 @@ def test_file_is_streaming_true_for_delta_methods(self): finally: os.unlink(temp_path) + def test_file_is_streaming_true_for_dmo_delta_methods(self): + temp_path = create_test_script(STREAMING_DMO_ENTRYPOINT) + try: + assert file_is_streaming(temp_path) is True + finally: + os.unlink(temp_path) + def test_file_is_streaming_false_for_batch(self): content = textwrap.dedent( """ @@ -778,13 +785,72 @@ def test_scan_file_streaming_dlo(self): finally: os.unlink(temp_path) - def test_scan_file_streaming_dmo_read(self): + def test_scan_file_streaming_dmo(self): temp_path = create_test_script(STREAMING_DMO_ENTRYPOINT) try: result = scan_file_streaming(temp_path) + assert result.read_dlo_deltas is False assert result.read_dmo_deltas is True assert result.read_layer == "dmo" - assert "Account_copy__dll" in result.write_dlo_deltas + assert "Account_copy__dlm" in result.write_dmo_deltas + assert result.write_layer == "dmo" + assert not result.write_dlo_deltas + finally: + os.unlink(temp_path) + + def test_scan_file_streaming_rejects_both_write_methods(self): + content = textwrap.dedent( + """ + from datacustomcode.client import StreamingClient + + client = StreamingClient() + deltas = client.read_dlo_deltas() + client.write_dlo_deltas("Account_copy__dll", deltas) + client.write_dmo_deltas("Account_copy__dlm", deltas) + """ + ) + temp_path = create_test_script(content) + try: + with pytest.raises(ValueError, match="Cannot write DLO and DMO deltas"): + scan_file_streaming(temp_path) + finally: + os.unlink(temp_path) + + def test_scan_file_streaming_rejects_dlo_read_dmo_write(self): + content = textwrap.dedent( + """ + from datacustomcode.client import StreamingClient + + client = StreamingClient() + deltas = client.read_dlo_deltas() + client.write_dmo_deltas("Account_copy__dlm", deltas) + """ + ) + temp_path = create_test_script(content) + try: + with pytest.raises( + ValueError, match=r"same layer \(got read=dlo, write=dmo\)" + ): + scan_file_streaming(temp_path) + finally: + os.unlink(temp_path) + + def test_scan_file_streaming_rejects_dmo_read_dlo_write(self): + content = textwrap.dedent( + """ + from datacustomcode.client import StreamingClient + + client = StreamingClient() + deltas = client.read_dmo_deltas() + client.write_dlo_deltas("Account_copy__dll", deltas) + """ + ) + temp_path = create_test_script(content) + try: + with pytest.raises( + ValueError, match=r"same layer \(got read=dmo, write=dlo\)" + ): + scan_file_streaming(temp_path) finally: os.unlink(temp_path) @@ -875,6 +941,7 @@ def test_reasserts_layer_from_code(self): assert updated["streamingSource"]["type"] == "dmo" assert updated["streamingSource"]["name"] == "Account__dlm" assert updated["permissions"]["read"]["dmo"] == ["Account__dlm"] + assert updated["permissions"]["write"] == {"dmo": ["Account_copy__dlm"]} def test_batch_entrypoint_drops_stale_streaming_source(self): """Switching a streaming entrypoint to batch clears streamingSource."""