streaming local test

src/datacustomcode/__init__.py

@@ -19,6 +19,8 @@
     "Credentials",
     "DefaultSparkEinsteinPredictions",
     "DefaultSparkLLMGateway",
+    "LocalDeltasReader",
+    "LocalDeltasWriter",
     "PrintDataCloudWriter",
     "QueryAPIDataCloudReader",
     "SparkEinsteinPredictions",
@@ -55,6 +57,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,37 @@
+# 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,184 @@
+# 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
+import sys
+from pathlib import Path
+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)
+        # Credentials captured here and forwarded to StreamingSourceSeeder
+        # on the first `run` against a source (see _try_seed).
+        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,
+        )
+        # Which delta method was called; the seeder's LIMIT 1 probe
+        # needs to target the right layer.
+        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, or run `datacustomcode init --use-in-feature "
+                "StreamingTransform` to get started."
+            )
+        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 empty.\n"
+                f"  Populate the stream source with at least one row in Data Cloud, "
+                f"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 additional "
+                f"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:
+        """Fetch schema + snapshot from the tenant for this source.
+
+        Called by _open_stream when _schema.json is missing (first run
+        against `name`, or the customer deleted the cache to force a
+        fresh seed). The `_` prefix is a Hadoop hidden-file convention;
+        Spark's FileStreamSource skips it so the schema cache doesn't
+        get delivered as a micro-batch alongside 000_seed.json.
+
+        Returns True when fixtures were written, False when the source
+        is empty (caller sys.exit(0)s with a friendly message). Query
+        API errors propagate as RuntimeError.
+        """
+        # Deferred import: batch-only runs never load the seeder module.
+        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,
+        )
+        # self._current_layer was set by read_dlo_deltas / read_dmo_deltas
+        # right before _open_stream — pass it through so the seeder's
+        # LIMIT 1 probe targets the correct read method.
+        return seeder.seed_source(
+            name, self._current_layer, str(self._fixtures_root)
+        )

src/datacustomcode/io/reader/streaming_seeder.py

@@ -0,0 +1,113 @@
+# 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 math
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, TYPE_CHECKING, Optional
+
+from datacustomcode.io import cdf
+from datacustomcode.io.reader.query_api import QueryAPIDataCloudReader
+
+if TYPE_CHECKING:
+    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
+        snapshot = 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