Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
10 changes: 10 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down Expand Up @@ -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.
Expand Down
19 changes: 19 additions & 0 deletions src/datacustomcode/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
8 changes: 1 addition & 7 deletions src/datacustomcode/deploy.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 "
Expand Down
34 changes: 33 additions & 1 deletion src/datacustomcode/io/writer/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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."
)
32 changes: 26 additions & 6 deletions src/datacustomcode/scan.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down Expand Up @@ -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:
Expand All @@ -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

Expand All @@ -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."""
Expand All @@ -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:
Expand Down Expand Up @@ -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]:
Expand All @@ -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:
Expand All @@ -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),
)


Expand Down Expand Up @@ -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},
}


Expand Down
6 changes: 6 additions & 0 deletions tests/io/writer/test_print.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
52 changes: 52 additions & 0 deletions tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand Down
20 changes: 9 additions & 11 deletions tests/test_deploy.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]),
),
)

Expand Down
Loading
Loading