support local testing of streaming transforms

src/datacustomcode/__init__.py

@@ -24,6 +24,8 @@
     "Credentials",
     "DefaultSparkEinsteinPredictions",
     "DefaultSparkLLMGateway",
+    "LocalDeltasReader",
+    "LocalDeltasWriter",
     "PrintDataCloudWriter",
     "QueryAPIDataCloudReader",
     "SparkEinsteinPredictions",
@@ -60,6 +62,14 @@ def __getattr__(name: str):
         from datacustomcode.io.reader.query_api import QueryAPIDataCloudReader
 
         return QueryAPIDataCloudReader
+    elif name == "LocalDeltasReader":
+        from datacustomcode.io.reader.local_deltas import LocalDeltasReader
+
+        return LocalDeltasReader
+    elif name == "LocalDeltasWriter":
+        from datacustomcode.io.writer.local_deltas import LocalDeltasWriter
+
+        return LocalDeltasWriter
     elif name == "SparkLLMGateway":
         from datacustomcode.llm_gateway import SparkLLMGateway
 

src/datacustomcode/io/cdf.py

@@ -0,0 +1,34 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import Final
+
+from datacustomcode.io.writer.base import MERGE_RECORD_TYPE_COLUMN, MergeRecordType
+
+COMMIT_VERSION: Final = "_commit_version"
+COMMIT_TIMESTAMP: Final = "_commit_timestamp"
+MERGE_RECORD_TYPE: Final = MERGE_RECORD_TYPE_COLUMN
+
+CDF_METADATA_COLUMNS: Final = (COMMIT_VERSION, COMMIT_TIMESTAMP, MERGE_RECORD_TYPE)
+
+__all__ = [
+    "COMMIT_VERSION",
+    "COMMIT_TIMESTAMP",
+    "MERGE_RECORD_TYPE",
+    "CDF_METADATA_COLUMNS",
+    "MergeRecordType",
+]

src/datacustomcode/io/reader/local_deltas.py

@@ -0,0 +1,164 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+import sys
+from typing import (
+    TYPE_CHECKING,
+    Optional,
+    Union,
+)
+
+from datacustomcode.config import config
+from datacustomcode.io import cdf
+from datacustomcode.io.reader.base import BaseDataCloudReader
+from datacustomcode.io.reader.query_api import QueryAPIDataCloudReader
+
+if TYPE_CHECKING:
+    from pyspark.sql import DataFrame as PySparkDataFrame, SparkSession
+    from pyspark.sql.types import AtomicType, StructType
+
+logger = logging.getLogger(__name__)
+
+
+class LocalDeltasReader(BaseDataCloudReader):
+    """Local reader for streaming transforms"""
+
+    CONFIG_NAME = "LocalDeltasReader"
+
+    def __init__(
+        self,
+        spark: SparkSession,
+        credentials_profile: str = "default",
+        dataspace: Optional[str] = None,
+        sf_cli_org: Optional[str] = None,
+        default_row_limit: Optional[int] = None,
+        fixtures_root: str = "payload/streaming_fixtures",
+    ) -> None:
+        super().__init__(spark)
+        self._fixtures_root = Path(fixtures_root)
+        self._credentials_profile = credentials_profile
+        self._dataspace = dataspace
+        self._sf_cli_org = sf_cli_org
+        # Reuse the batch reader for INITIAL_SYNC / REBUILD reads.
+        self._batch = QueryAPIDataCloudReader(
+            spark=spark,
+            credentials_profile=credentials_profile,
+            dataspace=dataspace,
+            sf_cli_org=sf_cli_org,
+            default_row_limit=default_row_limit,
+        )
+        self._current_layer = "dlo"
+
+    def read_dlo(
+        self,
+        name: str,
+        schema: Union[AtomicType, StructType, str, None] = None,
+    ) -> PySparkDataFrame:
+        return self._batch.read_dlo(name, schema)
+
+    def read_dmo(
+        self,
+        name: str,
+        schema: Union[AtomicType, StructType, str, None] = None,
+    ) -> PySparkDataFrame:
+        return self._batch.read_dmo(name, schema)
+
+    def read_dlo_deltas(self) -> PySparkDataFrame:
+        self._current_layer = "dlo"
+        return self._open_stream(self._streaming_source())
+
+    def read_dmo_deltas(self) -> PySparkDataFrame:
+        self._current_layer = "dmo"
+        return self._open_stream(self._streaming_source())
+
+    def _streaming_source(self) -> str:
+        source = config.streaming_source
+        if not source:
+            raise RuntimeError(
+                "No streaming source configured. Set streamingSource.name in "
+                "config.json."
+            )
+        return source
+
+    def _open_stream(self, name: str) -> PySparkDataFrame:
+        drop_dir = self._fixtures_root / name
+
+        if not (drop_dir / "_schema.json").exists():
+            seeded = self._try_seed(name)
+            if not seeded:
+                print(
+                    f"\nStreaming source {self._current_layer}='{name}' is "
+                    f"empty.\n"
+                    f"  Populate the stream source with at least one row in "
+                    f"Data Cloud, then re-run `datacustomcode run`.\n"
+                )
+                sys.exit(0)
+
+            print(
+                f"\nWrote sample streaming fixture in: {drop_dir} "
+                f"based on the contents of the streaming source {name}.\n"
+                f"  You may add more JSON files alongside it to simulate "
+                f"additional changes.\n"
+            )
+
+        schema = self._build_stream_schema(name)
+        return (
+            self.spark.readStream.format("json")
+            .schema(schema)
+            .option("maxFilesPerTrigger", 1)  # one file = one batch
+            .option("latestFirst", "false")  # oldest mtime first
+            .load(str(drop_dir))
+        )
+
+    def _build_stream_schema(self, name: str) -> "StructType":
+        """Compose (source schema + CDF metadata columns) for readStream."""
+        from pyspark.sql.types import (
+            LongType,
+            StringType,
+            StructField,
+            StructType,
+            TimestampType,
+        )
+
+        source_schema = self._resolve_source_schema(name)
+        cdf_fields = [
+            StructField(cdf.COMMIT_VERSION, LongType(), True),
+            StructField(cdf.COMMIT_TIMESTAMP, TimestampType(), True),
+            StructField(cdf.MERGE_RECORD_TYPE, StringType(), True),
+        ]
+        return StructType(list(source_schema.fields) + cdf_fields)
+
+    def _resolve_source_schema(self, name: str) -> "StructType":
+        from pyspark.sql.types import StructType
+
+        schema_file = self._fixtures_root / name / "_schema.json"
+        return StructType.fromJson(json.loads(schema_file.read_text()))
+
+    def _try_seed(self, name: str) -> bool:
+        """Creates source schema and sample change file."""
+        from datacustomcode.io.reader.streaming_seeder import StreamingSourceSeeder
+
+        seeder = StreamingSourceSeeder(
+            spark=self.spark,
+            credentials_profile=self._credentials_profile,
+            dataspace=self._dataspace,
+            sf_cli_org=self._sf_cli_org,
+        )
+        return seeder.seed_source(name, self._current_layer, str(self._fixtures_root))

src/datacustomcode/io/reader/streaming_seeder.py

@@ -0,0 +1,114 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+import json
+import math
+from pathlib import Path
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Optional,
+)
+
+from datacustomcode.io import cdf
+from datacustomcode.io.reader.query_api import QueryAPIDataCloudReader
+
+if TYPE_CHECKING:
+    import pandas as pd
+    from pyspark.sql import SparkSession
+
+SEED_LIMIT = 10
+
+
+def _clean_for_json(value: Any) -> Any:
+    if isinstance(value, float):
+        if math.isnan(value):
+            return None
+        if value.is_integer():
+            return int(value)
+    return value
+
+
+class StreamingSourceSeeder:
+    def __init__(
+        self,
+        spark: SparkSession,
+        credentials_profile: str = "default",
+        dataspace: Optional[str] = None,
+        sf_cli_org: Optional[str] = None,
+    ) -> None:
+        self.spark = spark
+        self.reader = QueryAPIDataCloudReader(
+            spark=spark,
+            credentials_profile=credentials_profile,
+            dataspace=dataspace,
+            sf_cli_org=sf_cli_org,
+        )
+
+    def seed_source(self, name: str, layer: str, fixtures_root: str) -> bool:
+        """Seed streaming fixtures and source schema.
+
+        Returns:
+            ``True`` when schema + fixtures are written
+            ``False`` when the source is empty.
+            Query API errors (missing creds, missing
+            source) propagate as ``RuntimeError``.
+        """
+        read_fn = self.reader.read_dlo if layer == "dlo" else self.reader.read_dmo
+
+        try:
+            head_df = read_fn(name).limit(1)
+        except Exception as exc:
+            raise RuntimeError(f"Failed to read {layer}='{name}': {exc}.") from exc
+
+        head_pandas = head_df.toPandas()
+        if len(head_pandas) == 0:
+            return False
+
+        schema = head_df.schema
+        out_dir = Path(fixtures_root) / name
+        out_dir.mkdir(parents=True, exist_ok=True)
+        schema_file_path = out_dir / "_schema.json"
+        schema_file_path.write_text(json.dumps(schema.jsonValue()))
+        print(
+            f"\nStreaming source schema written to {schema_file_path} "
+            f"({len(schema.fields)} fields)"
+        )
+
+        # Mix UPSERT and DELETE in the seed so the customer sees both
+        # operation types in the starter fixture. Annotated because
+        # pyspark's ``.toPandas()`` return type is ``PandasDataFrameLike``,
+        # which mypy doesn't recognize as having ``to_dict``.
+        snapshot: pd.DataFrame = read_fn(name).limit(SEED_LIMIT).toPandas()
+        seed_rows = []
+        for i, record in enumerate(snapshot.to_dict("records")):
+            op = cdf.MergeRecordType.DELETE if i % 2 else cdf.MergeRecordType.UPSERT
+
+            cleaned = {k: _clean_for_json(v) for k, v in record.items()}
+            seed_rows.append(
+                {
+                    **cleaned,
+                    cdf.COMMIT_VERSION: i + 1,
+                    cdf.COMMIT_TIMESTAMP: datetime.now(timezone.utc).isoformat(),
+                    cdf.MERGE_RECORD_TYPE: op.value,
+                }
+            )
+        (out_dir / "000_seed.json").write_text(
+            "\n".join(json.dumps(r, default=str) for r in seed_rows)
+        )
+        return True