streaming source from config

Author: jcatt-sf

src/datacustomcode/client.py

@@ -15,7 +15,6 @@
 from __future__ import annotations
 
 from enum import Enum
-import os
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -52,18 +51,22 @@
     from datacustomcode.spark.base import BaseSparkSessionProvider
 
 
-_STREAMING_SOURCE_ENV = "BYOC_STREAMING_SOURCE_NAME"
-
-
 def _streaming_source_name() -> str:
-    """Return the runtime streaming source name.
+    """Return the streaming transform's read-source name.
+
+    Resolved from ``config.streaming_source``, which ``run_entrypoint``
+    populates from config.json's ``permissions.read`` entry.
 
     Raises:
-        RuntimeError: If ``BYOC_STREAMING_SOURCE_NAME`` is not set
+        RuntimeError: If no ``streaming_source`` has been configured (e.g. the
+            transform's config.json has no ``permissions.read`` entry).
     """
-    source = os.environ.get(_STREAMING_SOURCE_ENV)
+    source = config.streaming_source
     if not source:
-        raise RuntimeError(f"{_STREAMING_SOURCE_ENV} is not set.")
+        raise RuntimeError(
+            "No streaming source configured. A streaming transform must declare "
+            "its read source in config.json under 'permissions.read'."
+        )
     return source
 
 
@@ -583,7 +586,7 @@ def write_dlo_deltas(
         """Write a streaming DataFrame of deltas to a DLO in Data Cloud.
 
         Starts a streaming query that writes each micro-batch to the
-        target DLO and returns the  ``StreamingQuery`` handle; the caller 
+        target DLO and returns the  ``StreamingQuery`` handle; the caller
         typically calls ``query.awaitTermination()``.
 
         Args:

src/datacustomcode/config.py

@@ -89,6 +89,9 @@ class ClientConfig(BaseConfig):
     spark_provider_config: Union[
         SparkProviderConfig[BaseSparkSessionProvider], None
     ] = None
+    # Source object name for a streaming (DELTA_SYNC) transform, populated by
+    # ``run_entrypoint`` from config.json's ``permissions.read``
+    streaming_source: Union[str, None] = None
 
     def update(self, other: ClientConfig) -> ClientConfig:
         """Merge this ClientConfig with another, respecting force flags.
@@ -116,6 +119,8 @@ def merge(
         self.spark_provider_config = merge(
             self.spark_provider_config, other.spark_provider_config
         )
+        if other.streaming_source is not None:
+            self.streaming_source = other.streaming_source
         return self
 
 

src/datacustomcode/run.py

@@ -42,6 +42,22 @@ def _set_config_option(config_obj, key: str, value: Optional[str]) -> None:
         config_obj.options[key] = value
 
 
+def _read_source_from_permissions(config_json: dict) -> Optional[str]:
+    """Return the read-source name from config.json ``permissions.read``.
+    """
+    permissions = config_json.get("permissions")
+    if not isinstance(permissions, dict):
+        return None
+    read = permissions.get("read")
+    if not isinstance(read, dict):
+        return None
+    for layer in ("dlo", "dmo"):
+        names = read.get(layer)
+        if names:
+            return names[0]
+    return None
+
+
 def _update_config_options(profile: Optional[str], sf_cli_org: Optional[str]):
     if sf_cli_org:
         config_key = "sf_cli_org"
@@ -125,6 +141,8 @@ def run_entrypoint(
         _set_config_option(config.reader_config, "dataspace", dataspace)
         _set_config_option(config.writer_config, "dataspace", dataspace)
 
+        config.streaming_source = _read_source_from_permissions(config_json)
+
     _update_config_options(profile, sf_cli_org)
 
     for dependency in dependencies:

tests/test_client.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import os
 from unittest.mock import MagicMock, patch
 
 from pyspark.sql import DataFrame, SparkSession
@@ -84,13 +83,17 @@ def mock_config(mock_spark):
 @pytest.fixture
 def reset_client():
     """Reset the client singletons (and the shared Spark session) between tests."""
+    from datacustomcode.client import config as client_config
+
     Client._instance = None
     StreamingClient._instance = None
     _BaseClient._shared_spark = None
+    client_config.streaming_source = None
     yield
     Client._instance = None
     StreamingClient._instance = None
     _BaseClient._shared_spark = None
+    client_config.streaming_source = None
 
 
 class TestClient:
@@ -298,29 +301,29 @@ def test_read_dlo_deltas(self, reset_client, mock_spark):
         reader.read_dlo_deltas.return_value = mock_df
 
         client = StreamingClient(reader=reader, writer=writer)
-        # The streaming source is resolved by the runtime and recorded from the
-        # env var it sets; the caller passes no name.
-        with patch.dict(
-            "os.environ", {"BYOC_STREAMING_SOURCE_NAME": "Account_std__dll"}
-        ):
+        
+        with patch("datacustomcode.client.config") as mock_config:
+            mock_config.streaming_source = "Account_std__dll"
             result = client.read_dlo_deltas()
 
         reader.read_dlo_deltas.assert_called_once_with()
         assert result is mock_df
         assert "Account_std__dll" in client._data_layer_history[DataCloudObjectType.DLO]
 
-    def test_read_dlo_deltas_without_source_env_raises(self, reset_client, mock_spark):
-        """Delta reads require the runtime source env var; absence fails fast."""
+    def test_read_dlo_deltas_without_configured_source_raises(
+        self, reset_client, mock_spark
+    ):
+        """Delta reads require a configured streaming source; absence fails fast."""
         reader = MagicMock(spec=BaseDataCloudReader)
         writer = MagicMock(spec=BaseDataCloudWriter)
 
         client = StreamingClient(reader=reader, writer=writer)
-        with patch.dict("os.environ", {}, clear=False):
-            os.environ.pop("BYOC_STREAMING_SOURCE_NAME", None)
+        with patch("datacustomcode.client.config") as mock_config:
+            mock_config.streaming_source = None
             with pytest.raises(RuntimeError) as exc_info:
                 client.read_dlo_deltas()
 
-        assert "BYOC_STREAMING_SOURCE_NAME" in str(exc_info.value)
+        assert "permissions.read" in str(exc_info.value)
         reader.read_dlo_deltas.assert_not_called()
 
     def test_read_dmo_deltas(self, reset_client, mock_spark):
@@ -330,9 +333,8 @@ def test_read_dmo_deltas(self, reset_client, mock_spark):
         reader.read_dmo_deltas.return_value = mock_df
 
         client = StreamingClient(reader=reader, writer=writer)
-        with patch.dict(
-            "os.environ", {"BYOC_STREAMING_SOURCE_NAME": "Account_model__dlm"}
-        ):
+        with patch("datacustomcode.client.config") as mock_config:
+            mock_config.streaming_source = "Account_model__dlm"
             result = client.read_dmo_deltas()
 
         reader.read_dmo_deltas.assert_called_once_with()
@@ -384,7 +386,8 @@ def test_streaming_read_write_flow(self, reset_client, mock_spark):
 
         client = StreamingClient(reader=reader, writer=writer)
 
-        with patch.dict("os.environ", {"BYOC_STREAMING_SOURCE_NAME": "source_dll"}):
+        with patch("datacustomcode.client.config") as mock_config:
+            mock_config.streaming_source = "source_dll"
             df = client.read_dlo_deltas()
         client.write_dlo_deltas("target_dll", df)
 

tests/test_run.py

@@ -488,3 +488,102 @@ def test_run_entrypoint_empty_dataspace_value(self):
                 os.unlink(entrypoint_file)
             if os.path.exists(config_json_path):
                 os.unlink(config_json_path)
+
+
+class TestReadSourceFromPermissions:
+    """`_read_source_from_permissions` extracts the streaming read source from
+    config.json's `permissions.read`."""
+
+    def test_returns_single_dlo(self):
+        from datacustomcode.run import _read_source_from_permissions
+
+        config_json = {"permissions": {"read": {"dlo": ["Account_std__dll"]}}}
+        assert _read_source_from_permissions(config_json) == "Account_std__dll"
+
+    def test_returns_single_dmo(self):
+        from datacustomcode.run import _read_source_from_permissions
+
+        config_json = {"permissions": {"read": {"dmo": ["Account_model__dlm"]}}}
+        assert _read_source_from_permissions(config_json) == "Account_model__dlm"
+
+    def test_dlo_preferred_when_both_present(self):
+        from datacustomcode.run import _read_source_from_permissions
+
+        config_json = {
+            "permissions": {"read": {"dlo": ["the_dll"], "dmo": ["the_dlm"]}}
+        }
+        assert _read_source_from_permissions(config_json) == "the_dll"
+
+    def test_returns_first_of_multiple(self):
+        from datacustomcode.run import _read_source_from_permissions
+
+        config_json = {"permissions": {"read": {"dlo": ["first__dll", "second__dll"]}}}
+        assert _read_source_from_permissions(config_json) == "first__dll"
+
+    @pytest.mark.parametrize(
+        "config_json",
+        [
+            {},
+            {"permissions": None},
+            {"permissions": {}},
+            {"permissions": {"read": None}},
+            {"permissions": {"read": {}}},
+            {"permissions": {"read": {"dlo": []}}},
+        ],
+    )
+    def test_returns_none_when_absent_or_empty(self, config_json):
+        from datacustomcode.run import _read_source_from_permissions
+
+        assert _read_source_from_permissions(config_json) is None
+
+
+class TestStreamingSourceScenarios:
+    """`run_entrypoint` populates `config.streaming_source` from config.json."""
+
+    def _run_capturing_streaming_source(self, config_json_body):
+        """Run an entrypoint that records config.streaming_source and return it."""
+        with tempfile.NamedTemporaryFile(suffix=".py", delete=False) as temp:
+            entrypoint_content = textwrap.dedent(
+                """
+                from datacustomcode.config import config
+                with open("streaming_source_output.txt", "w") as f:
+                    f.write(f"streaming_source: {config.streaming_source}")
+                """
+            )
+            temp.write(entrypoint_content.encode("utf-8"))
+            entrypoint_file = temp.name
+
+        entrypoint_dir = os.path.dirname(entrypoint_file)
+        config_json_path = os.path.join(entrypoint_dir, "config.json")
+        with open(config_json_path, "w") as f:
+            json.dump(config_json_body, f)
+
+        try:
+            run_entrypoint(
+                entrypoint=entrypoint_file,
+                config_file=None,
+                dependencies=[],
+                profile="default",
+            )
+            with open("streaming_source_output.txt", "r") as f:
+                return f.read()
+        finally:
+            if os.path.exists(entrypoint_file):
+                os.unlink(entrypoint_file)
+            if os.path.exists(config_json_path):
+                os.unlink(config_json_path)
+            if os.path.exists("streaming_source_output.txt"):
+                os.unlink("streaming_source_output.txt")
+
+    def test_streaming_source_set_from_permissions_read(self):
+        content = self._run_capturing_streaming_source(
+            {
+                "dataspace": "default",
+                "permissions": {"read": {"dlo": ["Account_std__dll"]}},
+            }
+        )
+        assert "streaming_source: Account_std__dll" in content
+
+    def test_streaming_source_none_for_batch_without_read(self):
+        content = self._run_capturing_streaming_source({"dataspace": "default"})
+        assert "streaming_source: None" in content