Skip to content

Commit 5f34e3e

Browse files
atulikumwenayoatulikumwenayo
authored
Merge pull request #150 from forcedotcom/localStreaming3
support local testing of streaming transforms
2 parents 12d4c97 + 340be1c commit 5f34e3e

11 files changed

Lines changed: 770 additions & 0 deletions

File tree

‎src/datacustomcode/__init__.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@
2424
"Credentials",
2525
"DefaultSparkEinsteinPredictions",
2626
"DefaultSparkLLMGateway",
27+
"LocalDeltasReader",
28+
"LocalDeltasWriter",
2729
"PrintDataCloudWriter",
2830
"QueryAPIDataCloudReader",
2931
"SparkEinsteinPredictions",
@@ -60,6 +62,14 @@ def __getattr__(name: str):
6062
from datacustomcode.io.reader.query_api import QueryAPIDataCloudReader
6163

6264
return QueryAPIDataCloudReader
65+
elif name == "LocalDeltasReader":
66+
from datacustomcode.io.reader.local_deltas import LocalDeltasReader
67+
68+
return LocalDeltasReader
69+
elif name == "LocalDeltasWriter":
70+
from datacustomcode.io.writer.local_deltas import LocalDeltasWriter
71+
72+
return LocalDeltasWriter
6373
elif name == "SparkLLMGateway":
6474
from datacustomcode.llm_gateway import SparkLLMGateway
6575

‎src/datacustomcode/io/cdf.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
# Copyright (c) 2025, Salesforce, Inc.
2+
# SPDX-License-Identifier: Apache-2
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
from __future__ import annotations
17+
18+
from typing import Final
19+
20+
from datacustomcode.io.writer.base import MERGE_RECORD_TYPE_COLUMN, MergeRecordType
21+
22+
COMMIT_VERSION: Final = "_commit_version"
23+
COMMIT_TIMESTAMP: Final = "_commit_timestamp"
24+
MERGE_RECORD_TYPE: Final = MERGE_RECORD_TYPE_COLUMN
25+
26+
CDF_METADATA_COLUMNS: Final = (COMMIT_VERSION, COMMIT_TIMESTAMP, MERGE_RECORD_TYPE)
27+
28+
__all__ = [
29+
"COMMIT_VERSION",
30+
"COMMIT_TIMESTAMP",
31+
"MERGE_RECORD_TYPE",
32+
"CDF_METADATA_COLUMNS",
33+
"MergeRecordType",
34+
]

‎src/datacustomcode/io/reader/local_deltas.py‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
# Copyright (c) 2025, Salesforce, Inc.
2+
# SPDX-License-Identifier: Apache-2
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
from __future__ import annotations
17+
18+
import json
19+
import logging
20+
from pathlib import Path
21+
import sys
22+
from typing import (
23+
TYPE_CHECKING,
24+
Optional,
25+
Union,
26+
)
27+
28+
from datacustomcode.config import config
29+
from datacustomcode.io import cdf
30+
from datacustomcode.io.reader.base import BaseDataCloudReader
31+
from datacustomcode.io.reader.query_api import QueryAPIDataCloudReader
32+
33+
if TYPE_CHECKING:
34+
from pyspark.sql import DataFrame as PySparkDataFrame, SparkSession
35+
from pyspark.sql.types import AtomicType, StructType
36+
37+
logger = logging.getLogger(__name__)
38+
39+
40+
class LocalDeltasReader(BaseDataCloudReader):
41+
"""Local reader for streaming transforms"""
42+
43+
CONFIG_NAME = "LocalDeltasReader"
44+
45+
def __init__(
46+
self,
47+
spark: SparkSession,
48+
credentials_profile: str = "default",
49+
dataspace: Optional[str] = None,
50+
sf_cli_org: Optional[str] = None,
51+
default_row_limit: Optional[int] = None,
52+
fixtures_root: str = "payload/streaming_fixtures",
53+
) -> None:
54+
super().__init__(spark)
55+
self._fixtures_root = Path(fixtures_root)
56+
self._credentials_profile = credentials_profile
57+
self._dataspace = dataspace
58+
self._sf_cli_org = sf_cli_org
59+
# Reuse the batch reader for INITIAL_SYNC / REBUILD reads.
60+
self._batch = QueryAPIDataCloudReader(
61+
spark=spark,
62+
credentials_profile=credentials_profile,
63+
dataspace=dataspace,
64+
sf_cli_org=sf_cli_org,
65+
default_row_limit=default_row_limit,
66+
)
67+
self._current_layer = "dlo"
68+
69+
def read_dlo(
70+
self,
71+
name: str,
72+
schema: Union[AtomicType, StructType, str, None] = None,
73+
) -> PySparkDataFrame:
74+
return self._batch.read_dlo(name, schema)
75+
76+
def read_dmo(
77+
self,
78+
name: str,
79+
schema: Union[AtomicType, StructType, str, None] = None,
80+
) -> PySparkDataFrame:
81+
return self._batch.read_dmo(name, schema)
82+
83+
def read_dlo_deltas(self) -> PySparkDataFrame:
84+
self._current_layer = "dlo"
85+
return self._open_stream(self._streaming_source())
86+
87+
def read_dmo_deltas(self) -> PySparkDataFrame:
88+
self._current_layer = "dmo"
89+
return self._open_stream(self._streaming_source())
90+
91+
def _streaming_source(self) -> str:
92+
source = config.streaming_source
93+
if not source:
94+
raise RuntimeError(
95+
"No streaming source configured. Set streamingSource.name in "
96+
"config.json."
97+
)
98+
return source
99+
100+
def _open_stream(self, name: str) -> PySparkDataFrame:
101+
drop_dir = self._fixtures_root / name
102+
103+
if not (drop_dir / "_schema.json").exists():
104+
seeded = self._try_seed(name)
105+
if not seeded:
106+
print(
107+
f"\nStreaming source {self._current_layer}='{name}' is "
108+
f"empty.\n"
109+
f" Populate the stream source with at least one row in "
110+
f"Data Cloud, then re-run `datacustomcode run`.\n"
111+
)
112+
sys.exit(0)
113+
114+
print(
115+
f"\nWrote sample streaming fixture in: {drop_dir} "
116+
f"based on the contents of the streaming source {name}.\n"
117+
f" You may add more JSON files alongside it to simulate "
118+
f"additional changes.\n"
119+
)
120+
121+
schema = self._build_stream_schema(name)
122+
return (
123+
self.spark.readStream.format("json")
124+
.schema(schema)
125+
.option("maxFilesPerTrigger", 1) # one file = one batch
126+
.option("latestFirst", "false") # oldest mtime first
127+
.load(str(drop_dir))
128+
)
129+
130+
def _build_stream_schema(self, name: str) -> "StructType":
131+
"""Compose (source schema + CDF metadata columns) for readStream."""
132+
from pyspark.sql.types import (
133+
LongType,
134+
StringType,
135+
StructField,
136+
StructType,
137+
TimestampType,
138+
)
139+
140+
source_schema = self._resolve_source_schema(name)
141+
cdf_fields = [
142+
StructField(cdf.COMMIT_VERSION, LongType(), True),
143+
StructField(cdf.COMMIT_TIMESTAMP, TimestampType(), True),
144+
StructField(cdf.MERGE_RECORD_TYPE, StringType(), True),
145+
]
146+
return StructType(list(source_schema.fields) + cdf_fields)
147+
148+
def _resolve_source_schema(self, name: str) -> "StructType":
149+
from pyspark.sql.types import StructType
150+
151+
schema_file = self._fixtures_root / name / "_schema.json"
152+
return StructType.fromJson(json.loads(schema_file.read_text()))
153+
154+
def _try_seed(self, name: str) -> bool:
155+
"""Creates source schema and sample change file."""
156+
from datacustomcode.io.reader.streaming_seeder import StreamingSourceSeeder
157+
158+
seeder = StreamingSourceSeeder(
159+
spark=self.spark,
160+
credentials_profile=self._credentials_profile,
161+
dataspace=self._dataspace,
162+
sf_cli_org=self._sf_cli_org,
163+
)
164+
return seeder.seed_source(name, self._current_layer, str(self._fixtures_root))

‎src/datacustomcode/io/reader/streaming_seeder.py‎

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
# Copyright (c) 2025, Salesforce, Inc.
2+
# SPDX-License-Identifier: Apache-2
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
from __future__ import annotations
17+
18+
from datetime import datetime, timezone
19+
import json
20+
import math
21+
from pathlib import Path
22+
from typing import (
23+
TYPE_CHECKING,
24+
Any,
25+
Optional,
26+
)
27+
28+
from datacustomcode.io import cdf
29+
from datacustomcode.io.reader.query_api import QueryAPIDataCloudReader
30+
31+
if TYPE_CHECKING:
32+
import pandas as pd
33+
from pyspark.sql import SparkSession
34+
35+
SEED_LIMIT = 10
36+
37+
38+
def _clean_for_json(value: Any) -> Any:
39+
if isinstance(value, float):
40+
if math.isnan(value):
41+
return None
42+
if value.is_integer():
43+
return int(value)
44+
return value
45+
46+
47+
class StreamingSourceSeeder:
48+
def __init__(
49+
self,
50+
spark: SparkSession,
51+
credentials_profile: str = "default",
52+
dataspace: Optional[str] = None,
53+
sf_cli_org: Optional[str] = None,
54+
) -> None:
55+
self.spark = spark
56+
self.reader = QueryAPIDataCloudReader(
57+
spark=spark,
58+
credentials_profile=credentials_profile,
59+
dataspace=dataspace,
60+
sf_cli_org=sf_cli_org,
61+
)
62+
63+
def seed_source(self, name: str, layer: str, fixtures_root: str) -> bool:
64+
"""Seed streaming fixtures and source schema.
65+
66+
Returns:
67+
``True`` when schema + fixtures are written
68+
``False`` when the source is empty.
69+
Query API errors (missing creds, missing
70+
source) propagate as ``RuntimeError``.
71+
"""
72+
read_fn = self.reader.read_dlo if layer == "dlo" else self.reader.read_dmo
73+
74+
try:
75+
head_df = read_fn(name).limit(1)
76+
except Exception as exc:
77+
raise RuntimeError(f"Failed to read {layer}='{name}': {exc}.") from exc
78+
79+
head_pandas = head_df.toPandas()
80+
if len(head_pandas) == 0:
81+
return False
82+
83+
schema = head_df.schema
84+
out_dir = Path(fixtures_root) / name
85+
out_dir.mkdir(parents=True, exist_ok=True)
86+
schema_file_path = out_dir / "_schema.json"
87+
schema_file_path.write_text(json.dumps(schema.jsonValue()))
88+
print(
89+
f"\nStreaming source schema written to {schema_file_path} "
90+
f"({len(schema.fields)} fields)"
91+
)
92+
93+
# Mix UPSERT and DELETE in the seed so the customer sees both
94+
# operation types in the starter fixture. Annotated because
95+
# pyspark's ``.toPandas()`` return type is ``PandasDataFrameLike``,
96+
# which mypy doesn't recognize as having ``to_dict``.
97+
snapshot: pd.DataFrame = read_fn(name).limit(SEED_LIMIT).toPandas()
98+
seed_rows = []
99+
for i, record in enumerate(snapshot.to_dict("records")):
100+
op = cdf.MergeRecordType.DELETE if i % 2 else cdf.MergeRecordType.UPSERT
101+
102+
cleaned = {k: _clean_for_json(v) for k, v in record.items()}
103+
seed_rows.append(
104+
{
105+
**cleaned,
106+
cdf.COMMIT_VERSION: i + 1,
107+
cdf.COMMIT_TIMESTAMP: datetime.now(timezone.utc).isoformat(),
108+
cdf.MERGE_RECORD_TYPE: op.value,
109+
}
110+
)
111+
(out_dir / "000_seed.json").write_text(
112+
"\n".join(json.dumps(r, default=str) for r in seed_rows)
113+
)
114+
return True

0 commit comments

Comments
 (0)