Merge pull request #163 from forcedotcom/feat/write-dmo-delta-support

Add write_dmo_deltas to StreamingClient

Author: andrewmiller-sf

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

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.

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.

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 "

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."
+        )

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},
     }
 
 

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"

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)

tests/test_deploy.py

@@ -1683,34 +1683,32 @@ 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",
             dataspace="default",
             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"]),
                 ),
             )