salesforce-data-customcode 6.0.8.dev1__py3-none-any.whl → 6.1.0.dev1__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- datacustomcode/__init__.py +5 -0
- datacustomcode/cli.py +29 -4
- datacustomcode/client.py +211 -100
- datacustomcode/config.py +5 -0
- datacustomcode/constants.py +8 -0
- datacustomcode/deploy.py +55 -16
- datacustomcode/einstein_predictions/spark_default.py +88 -27
- datacustomcode/io/reader/base.py +42 -0
- datacustomcode/io/writer/base.py +33 -0
- datacustomcode/run.py +11 -0
- datacustomcode/scan.py +164 -29
- datacustomcode/template.py +13 -1
- datacustomcode/templates/script/examples/streaming_deltas/entrypoint.py +49 -0
- {salesforce_data_customcode-6.0.8.dev1.dist-info → salesforce_data_customcode-6.1.0.dev1.dist-info}/METADATA +40 -3
- {salesforce_data_customcode-6.0.8.dev1.dist-info → salesforce_data_customcode-6.1.0.dev1.dist-info}/RECORD +18 -17
- {salesforce_data_customcode-6.0.8.dev1.dist-info → salesforce_data_customcode-6.1.0.dev1.dist-info}/WHEEL +0 -0
- {salesforce_data_customcode-6.0.8.dev1.dist-info → salesforce_data_customcode-6.1.0.dev1.dist-info}/entry_points.txt +0 -0
- {salesforce_data_customcode-6.0.8.dev1.dist-info → salesforce_data_customcode-6.1.0.dev1.dist-info}/licenses/LICENSE.txt +0 -0
datacustomcode/__init__.py
CHANGED
|
@@ -23,6 +23,7 @@ __all__ = [
|
|
|
23
23
|
"QueryAPIDataCloudReader",
|
|
24
24
|
"SparkEinsteinPredictions",
|
|
25
25
|
"SparkLLMGateway",
|
|
26
|
+
"StreamingClient",
|
|
26
27
|
"einstein_predict_col",
|
|
27
28
|
"llm_gateway_generate_text_col",
|
|
28
29
|
]
|
|
@@ -34,6 +35,10 @@ def __getattr__(name: str):
|
|
|
34
35
|
from datacustomcode.client import Client
|
|
35
36
|
|
|
36
37
|
return Client
|
|
38
|
+
elif name == "StreamingClient":
|
|
39
|
+
from datacustomcode.client import StreamingClient
|
|
40
|
+
|
|
41
|
+
return StreamingClient
|
|
37
42
|
elif name == "AuthType":
|
|
38
43
|
from datacustomcode.credentials import AuthType
|
|
39
44
|
|
datacustomcode/cli.py
CHANGED
|
@@ -283,10 +283,19 @@ def deploy(
|
|
|
283
283
|
)
|
|
284
284
|
@click.option(
|
|
285
285
|
"--use-in-feature",
|
|
286
|
-
|
|
287
|
-
|
|
286
|
+
"-u",
|
|
287
|
+
default=None,
|
|
288
|
+
help=(
|
|
289
|
+
"Invoke option for this package. For scripts: 'BatchTransform' "
|
|
290
|
+
"(default) or 'StreamingTransform'. For functions: 'SearchIndexChunking'."
|
|
291
|
+
),
|
|
288
292
|
)
|
|
289
293
|
def init(directory: str, code_type: str, use_in_feature: Optional[str]):
|
|
294
|
+
from datacustomcode.constants import (
|
|
295
|
+
SCRIPT_USE_IN_FEATURE_BATCH,
|
|
296
|
+
SCRIPT_USE_IN_FEATURE_OPTIONS,
|
|
297
|
+
SCRIPT_USE_IN_FEATURE_STREAMING,
|
|
298
|
+
)
|
|
290
299
|
from datacustomcode.scan import (
|
|
291
300
|
dc_config_json_from_file,
|
|
292
301
|
update_config,
|
|
@@ -294,9 +303,23 @@ def init(directory: str, code_type: str, use_in_feature: Optional[str]):
|
|
|
294
303
|
)
|
|
295
304
|
from datacustomcode.template import copy_function_template, copy_script_template
|
|
296
305
|
|
|
306
|
+
streaming = False
|
|
307
|
+
if code_type == "script":
|
|
308
|
+
use_in_feature = use_in_feature or SCRIPT_USE_IN_FEATURE_BATCH
|
|
309
|
+
if use_in_feature not in SCRIPT_USE_IN_FEATURE_OPTIONS:
|
|
310
|
+
click.secho(
|
|
311
|
+
f"Error: Invalid --use-in-feature '{use_in_feature}' for a "
|
|
312
|
+
f"script. Valid options: {', '.join(SCRIPT_USE_IN_FEATURE_OPTIONS)}.",
|
|
313
|
+
fg="red",
|
|
314
|
+
)
|
|
315
|
+
raise click.Abort()
|
|
316
|
+
streaming = use_in_feature == SCRIPT_USE_IN_FEATURE_STREAMING
|
|
317
|
+
else:
|
|
318
|
+
use_in_feature = use_in_feature or "SearchIndexChunking"
|
|
319
|
+
|
|
297
320
|
click.echo("Copying template to " + click.style(directory, fg="blue", bold=True))
|
|
298
321
|
if code_type == "script":
|
|
299
|
-
copy_script_template(directory)
|
|
322
|
+
copy_script_template(directory, streaming=streaming)
|
|
300
323
|
elif code_type == "function":
|
|
301
324
|
copy_function_template(directory, use_in_feature)
|
|
302
325
|
entrypoint_path = os.path.join(directory, PAYLOAD_DIR, ENTRYPOINT_FILE)
|
|
@@ -306,7 +329,9 @@ def init(directory: str, code_type: str, use_in_feature: Optional[str]):
|
|
|
306
329
|
sdk_config = {"type": code_type}
|
|
307
330
|
write_sdk_config(directory, sdk_config)
|
|
308
331
|
|
|
309
|
-
config_json = dc_config_json_from_file(
|
|
332
|
+
config_json = dc_config_json_from_file(
|
|
333
|
+
entrypoint_path, code_type, streaming=streaming
|
|
334
|
+
)
|
|
310
335
|
with open(config_location, "w") as f:
|
|
311
336
|
json.dump(config_json, f, indent=2)
|
|
312
337
|
|
datacustomcode/client.py
CHANGED
|
@@ -21,7 +21,9 @@ from typing import (
|
|
|
21
21
|
ClassVar,
|
|
22
22
|
Dict,
|
|
23
23
|
Optional,
|
|
24
|
+
TypeVar,
|
|
24
25
|
Union,
|
|
26
|
+
cast,
|
|
25
27
|
)
|
|
26
28
|
|
|
27
29
|
from datacustomcode.config import config
|
|
@@ -34,7 +36,12 @@ from datacustomcode.spark.default import DefaultSparkSessionProvider
|
|
|
34
36
|
if TYPE_CHECKING:
|
|
35
37
|
from pathlib import Path
|
|
36
38
|
|
|
37
|
-
from pyspark.sql import
|
|
39
|
+
from pyspark.sql import (
|
|
40
|
+
Column,
|
|
41
|
+
DataFrame as PySparkDataFrame,
|
|
42
|
+
SparkSession,
|
|
43
|
+
)
|
|
44
|
+
from pyspark.sql.streaming import StreamingQuery
|
|
38
45
|
|
|
39
46
|
from datacustomcode.einstein_predictions.spark_base import SparkEinsteinPredictions
|
|
40
47
|
from datacustomcode.einstein_predictions.types import PredictionType
|
|
@@ -44,6 +51,40 @@ if TYPE_CHECKING:
|
|
|
44
51
|
from datacustomcode.spark.base import BaseSparkSessionProvider
|
|
45
52
|
|
|
46
53
|
|
|
54
|
+
def _streaming_source_name() -> str:
|
|
55
|
+
"""Return the streaming transform's read-source name.
|
|
56
|
+
|
|
57
|
+
Resolved from ``config.streaming_source``, which ``run_entrypoint``
|
|
58
|
+
populates from config.json's ``streamingSource`` field.
|
|
59
|
+
|
|
60
|
+
Raises:
|
|
61
|
+
RuntimeError: If no ``streaming_source`` has been configured (e.g. the
|
|
62
|
+
transform's config.json has no ``streamingSource`` field).
|
|
63
|
+
"""
|
|
64
|
+
source = config.streaming_source
|
|
65
|
+
if not source:
|
|
66
|
+
raise RuntimeError(
|
|
67
|
+
"No streaming source configured. A streaming transform must declare "
|
|
68
|
+
"its read source in config.json under 'streamingSource'."
|
|
69
|
+
)
|
|
70
|
+
return source
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
def _active_client() -> "_BaseClient":
|
|
74
|
+
"""Return the client backing the module-level Spark column helpers.
|
|
75
|
+
|
|
76
|
+
Prefers an already-initialized singleton so a streaming job reuses its
|
|
77
|
+
:class:`StreamingClient` (and a batch job its :class:`Client`) rather than
|
|
78
|
+
forcing an unrelated client into existence. Falls back to building the
|
|
79
|
+
batch :class:`Client` when neither has been created yet.
|
|
80
|
+
"""
|
|
81
|
+
if Client._instance is not None:
|
|
82
|
+
return Client._instance
|
|
83
|
+
if StreamingClient._instance is not None:
|
|
84
|
+
return StreamingClient._instance
|
|
85
|
+
return Client()
|
|
86
|
+
|
|
87
|
+
|
|
47
88
|
def _build_spark_llm_gateway() -> "SparkLLMGateway":
|
|
48
89
|
"""Instantiate the SDK-configured :class:`SparkLLMGateway`.
|
|
49
90
|
|
|
@@ -99,7 +140,7 @@ def llm_gateway_generate_text_col(
|
|
|
99
140
|
the generated text; on failure, ``status == "ERROR"`` and the
|
|
100
141
|
``error_*`` fields carry diagnostic detail.
|
|
101
142
|
"""
|
|
102
|
-
gateway =
|
|
143
|
+
gateway = _active_client()._get_spark_llm_gateway()
|
|
103
144
|
return gateway.llm_gateway_generate_text_col(template, values, model_id=model_id)
|
|
104
145
|
|
|
105
146
|
|
|
@@ -161,7 +202,7 @@ def einstein_predict_col(
|
|
|
161
202
|
the JSON-serialized prediction payload; on failure, ``status ==
|
|
162
203
|
"ERROR"`` and the ``error_*`` fields carry diagnostic detail.
|
|
163
204
|
"""
|
|
164
|
-
predictions =
|
|
205
|
+
predictions = _active_client()._get_spark_einstein_predictions()
|
|
165
206
|
return predictions.einstein_predict_col(
|
|
166
207
|
model_api_name, prediction_type, features, settings=settings
|
|
167
208
|
)
|
|
@@ -205,39 +246,39 @@ class DataCloudAccessLayerException(Exception):
|
|
|
205
246
|
return msg
|
|
206
247
|
|
|
207
248
|
|
|
208
|
-
|
|
209
|
-
|
|
249
|
+
_ClientT = TypeVar("_ClientT", bound="_BaseClient")
|
|
250
|
+
|
|
251
|
+
|
|
252
|
+
class _BaseClient:
|
|
253
|
+
"""Shared machinery for the Data Cloud client singletons.
|
|
210
254
|
|
|
211
|
-
|
|
212
|
-
|
|
213
|
-
can read from DLOs and write to DLOs, read from DMOs and write to DMOs, but you
|
|
214
|
-
cannot read from DLOs and write to DMOs or read from DMOs and write to DLOs.
|
|
215
|
-
Furthermore you cannot mix during merging tables. This class is a singleton to
|
|
216
|
-
prevent accidental mixing of DLOs and DMOs.
|
|
255
|
+
Holds the wiring common to :class:`Client` (batch) and
|
|
256
|
+
:class:`StreamingClient`
|
|
217
257
|
|
|
218
|
-
|
|
219
|
-
|
|
220
|
-
behavior once deployed to Data Cloud. By default, the client intercepts all
|
|
221
|
-
read/write operations and mocks access to Data Cloud. For example, during
|
|
222
|
-
writing, we print to the console instead of writing to Data Cloud.
|
|
258
|
+
This base class is not meant to be instantiated directly; use
|
|
259
|
+
:class:`Client` or :class:`StreamingClient`.
|
|
223
260
|
|
|
224
261
|
Args:
|
|
225
|
-
finder: Find a file path
|
|
226
262
|
reader: A custom reader to use for reading Data Cloud objects.
|
|
227
263
|
writer: A custom writer to use for writing Data Cloud objects.
|
|
264
|
+
spark_provider: Optional custom :class:`BaseSparkSessionProvider`.
|
|
228
265
|
spark_llm_gateway: Optional custom :class:`SparkLLMGateway`.
|
|
229
266
|
spark_einstein_predictions: Optional custom
|
|
230
267
|
:class:`SparkEinsteinPredictions`.
|
|
231
|
-
|
|
232
|
-
Example:
|
|
233
|
-
>>> client = Client()
|
|
234
|
-
>>> file_path = client.find_file_path("data.csv")
|
|
235
|
-
>>> dlo = client.read_dlo("my_dlo")
|
|
236
|
-
>>> client.write_to_dmo("my_dmo", dlo)
|
|
237
|
-
>>> answer = client.llm_gateway_generate_text("Generate a greeting message")
|
|
238
268
|
"""
|
|
239
269
|
|
|
240
|
-
_instance:
|
|
270
|
+
# Each concrete subclass gets its own ``_instance`` slot: reads fall through
|
|
271
|
+
# to this base default of ``None``, but ``cls._instance = ...`` in __new__
|
|
272
|
+
# always writes to the subclass, so ``Client`` and ``StreamingClient`` never
|
|
273
|
+
# share an instance.
|
|
274
|
+
_instance: ClassVar[Optional[_BaseClient]] = None
|
|
275
|
+
# Process-wide Spark session shared across BOTH client types. Unlike
|
|
276
|
+
# ``_instance``, this is written via ``_BaseClient._shared_spark`` (never
|
|
277
|
+
# ``cls._shared_spark``), so the slot lives on the base class and a
|
|
278
|
+
# ``Client`` and a ``StreamingClient`` in the same process reuse one session
|
|
279
|
+
# — and therefore one underlying connection — instead of opening two
|
|
280
|
+
# containing differing state
|
|
281
|
+
_shared_spark: ClassVar[Optional[SparkSession]] = None
|
|
241
282
|
_reader: BaseDataCloudReader
|
|
242
283
|
_writer: BaseDataCloudWriter
|
|
243
284
|
_file: DefaultFindFilePath
|
|
@@ -247,37 +288,44 @@ class Client:
|
|
|
247
288
|
_code_type: str
|
|
248
289
|
|
|
249
290
|
def __new__(
|
|
250
|
-
cls,
|
|
291
|
+
cls: type[_ClientT],
|
|
251
292
|
reader: Optional[BaseDataCloudReader] = None,
|
|
252
293
|
writer: Optional[BaseDataCloudWriter] = None,
|
|
253
294
|
spark_provider: Optional[BaseSparkSessionProvider] = None,
|
|
254
295
|
spark_llm_gateway: Optional[SparkLLMGateway] = None,
|
|
255
296
|
spark_einstein_predictions: Optional[SparkEinsteinPredictions] = None,
|
|
256
297
|
code_type: str = "script",
|
|
257
|
-
) ->
|
|
298
|
+
) -> _ClientT:
|
|
258
299
|
|
|
259
300
|
if cls._instance is None:
|
|
260
|
-
|
|
261
|
-
|
|
262
|
-
|
|
301
|
+
instance = super().__new__(cls)
|
|
302
|
+
instance._spark_llm_gateway = spark_llm_gateway
|
|
303
|
+
instance._spark_einstein_predictions = spark_einstein_predictions
|
|
263
304
|
# Initialize Readers and Writers from config
|
|
264
305
|
# and/or provided reader and writer
|
|
265
306
|
if reader is None or writer is None:
|
|
266
|
-
# We need a spark because we will initialize readers and writers
|
|
267
|
-
if
|
|
268
|
-
|
|
269
|
-
|
|
270
|
-
|
|
271
|
-
|
|
272
|
-
provider: BaseSparkSessionProvider
|
|
273
|
-
if spark_provider is not None:
|
|
274
|
-
provider = spark_provider
|
|
275
|
-
elif config.spark_provider_config is not None:
|
|
276
|
-
provider = config.spark_provider_config.to_object()
|
|
307
|
+
# We need a spark because we will initialize readers and writers.
|
|
308
|
+
# Reuse the process-wide session if one client already built it,
|
|
309
|
+
# so a Client and a StreamingClient share a single connection.
|
|
310
|
+
if _BaseClient._shared_spark is not None:
|
|
311
|
+
spark = _BaseClient._shared_spark
|
|
277
312
|
else:
|
|
278
|
-
|
|
279
|
-
|
|
280
|
-
|
|
313
|
+
if config.spark_config is None:
|
|
314
|
+
raise ValueError(
|
|
315
|
+
"Spark config is required when reader/writer is not "
|
|
316
|
+
"provided"
|
|
317
|
+
)
|
|
318
|
+
|
|
319
|
+
provider: BaseSparkSessionProvider
|
|
320
|
+
if spark_provider is not None:
|
|
321
|
+
provider = spark_provider
|
|
322
|
+
elif config.spark_provider_config is not None:
|
|
323
|
+
provider = config.spark_provider_config.to_object()
|
|
324
|
+
else:
|
|
325
|
+
provider = DefaultSparkSessionProvider()
|
|
326
|
+
|
|
327
|
+
spark = provider.get_session(config.spark_config)
|
|
328
|
+
_BaseClient._shared_spark = spark
|
|
281
329
|
|
|
282
330
|
if config.reader_config is None and reader is None:
|
|
283
331
|
raise ValueError(
|
|
@@ -300,66 +348,17 @@ class Client:
|
|
|
300
348
|
else:
|
|
301
349
|
writer_init = writer
|
|
302
350
|
|
|
303
|
-
|
|
304
|
-
|
|
305
|
-
|
|
306
|
-
|
|
351
|
+
instance._reader = reader_init
|
|
352
|
+
instance._writer = writer_init
|
|
353
|
+
instance._file = DefaultFindFilePath()
|
|
354
|
+
instance._data_layer_history = {
|
|
307
355
|
DataCloudObjectType.DLO: set(),
|
|
308
356
|
DataCloudObjectType.DMO: set(),
|
|
309
357
|
}
|
|
310
|
-
|
|
358
|
+
cls._instance = instance
|
|
359
|
+
elif reader is not None or writer is not None:
|
|
311
360
|
raise ValueError("Cannot set reader or writer after client is initialized")
|
|
312
|
-
return cls._instance
|
|
313
|
-
|
|
314
|
-
def read_dlo(self, name: str) -> PySparkDataFrame:
|
|
315
|
-
"""Read a DLO from Data Cloud.
|
|
316
|
-
|
|
317
|
-
Args:
|
|
318
|
-
name: The name of the DLO to read.
|
|
319
|
-
|
|
320
|
-
Returns:
|
|
321
|
-
A PySpark DataFrame containing the DLO data.
|
|
322
|
-
"""
|
|
323
|
-
self._record_dlo_access(name)
|
|
324
|
-
return self._reader.read_dlo(name) # type: ignore[no-any-return]
|
|
325
|
-
|
|
326
|
-
def read_dmo(self, name: str) -> PySparkDataFrame:
|
|
327
|
-
"""Read a DMO from Data Cloud.
|
|
328
|
-
|
|
329
|
-
Args:
|
|
330
|
-
name: The name of the DMO to read.
|
|
331
|
-
|
|
332
|
-
Returns:
|
|
333
|
-
A PySpark DataFrame containing the DMO data.
|
|
334
|
-
"""
|
|
335
|
-
self._record_dmo_access(name)
|
|
336
|
-
return self._reader.read_dmo(name) # type: ignore[no-any-return]
|
|
337
|
-
|
|
338
|
-
def write_to_dlo(
|
|
339
|
-
self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs
|
|
340
|
-
) -> None:
|
|
341
|
-
"""Write a PySpark DataFrame to a DLO in Data Cloud.
|
|
342
|
-
|
|
343
|
-
Args:
|
|
344
|
-
name: The name of the DLO to write to.
|
|
345
|
-
dataframe: The PySpark DataFrame to write.
|
|
346
|
-
write_mode: The write mode to use for writing to the DLO.
|
|
347
|
-
"""
|
|
348
|
-
self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO)
|
|
349
|
-
return self._writer.write_to_dlo(name, dataframe, write_mode, **kwargs) # type: ignore[no-any-return]
|
|
350
|
-
|
|
351
|
-
def write_to_dmo(
|
|
352
|
-
self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs
|
|
353
|
-
) -> None:
|
|
354
|
-
"""Write a PySpark DataFrame to a DMO in Data Cloud.
|
|
355
|
-
|
|
356
|
-
Args:
|
|
357
|
-
name: The name of the DMO to write to.
|
|
358
|
-
dataframe: The PySpark DataFrame to write.
|
|
359
|
-
write_mode: The write mode to use for writing to the DMO.
|
|
360
|
-
"""
|
|
361
|
-
self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DLO)
|
|
362
|
-
return self._writer.write_to_dmo(name, dataframe, write_mode, **kwargs) # type: ignore[no-any-return]
|
|
361
|
+
return cast(_ClientT, cls._instance)
|
|
363
362
|
|
|
364
363
|
def find_file_path(self, file_name: str) -> Path:
|
|
365
364
|
"""Resolve a bundled file shipped in the package to an absolute path.
|
|
@@ -487,3 +486,115 @@ class Client:
|
|
|
487
486
|
|
|
488
487
|
def _record_dmo_access(self, name: str) -> None:
|
|
489
488
|
self._data_layer_history[DataCloudObjectType.DMO].add(name)
|
|
489
|
+
|
|
490
|
+
|
|
491
|
+
class Client(_BaseClient):
|
|
492
|
+
"""Entrypoint for batch access to Data Cloud objects.
|
|
493
|
+
|
|
494
|
+
This is the object used to read and write bounded snapshots of Data Cloud
|
|
495
|
+
DLOs and DMOs.
|
|
496
|
+
"""
|
|
497
|
+
|
|
498
|
+
_instance: ClassVar[Optional[Client]] = None
|
|
499
|
+
|
|
500
|
+
def read_dlo(self, name: str) -> PySparkDataFrame:
|
|
501
|
+
"""Read a DLO from Data Cloud.
|
|
502
|
+
|
|
503
|
+
Args:
|
|
504
|
+
name: The name of the DLO to read.
|
|
505
|
+
|
|
506
|
+
Returns:
|
|
507
|
+
A PySpark DataFrame containing the DLO data.
|
|
508
|
+
"""
|
|
509
|
+
self._record_dlo_access(name)
|
|
510
|
+
return self._reader.read_dlo(name) # type: ignore[no-any-return]
|
|
511
|
+
|
|
512
|
+
def read_dmo(self, name: str) -> PySparkDataFrame:
|
|
513
|
+
"""Read a DMO from Data Cloud.
|
|
514
|
+
|
|
515
|
+
Args:
|
|
516
|
+
name: The name of the DMO to read.
|
|
517
|
+
|
|
518
|
+
Returns:
|
|
519
|
+
A PySpark DataFrame containing the DMO data.
|
|
520
|
+
"""
|
|
521
|
+
self._record_dmo_access(name)
|
|
522
|
+
return self._reader.read_dmo(name) # type: ignore[no-any-return]
|
|
523
|
+
|
|
524
|
+
def write_to_dlo(
|
|
525
|
+
self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs
|
|
526
|
+
) -> None:
|
|
527
|
+
"""Write a PySpark DataFrame to a DLO in Data Cloud.
|
|
528
|
+
|
|
529
|
+
Args:
|
|
530
|
+
name: The name of the DLO to write to.
|
|
531
|
+
dataframe: The PySpark DataFrame to write.
|
|
532
|
+
write_mode: The write mode to use for writing to the DLO.
|
|
533
|
+
"""
|
|
534
|
+
self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO)
|
|
535
|
+
return self._writer.write_to_dlo(name, dataframe, write_mode, **kwargs) # type: ignore[no-any-return]
|
|
536
|
+
|
|
537
|
+
def write_to_dmo(
|
|
538
|
+
self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs
|
|
539
|
+
) -> None:
|
|
540
|
+
"""Write a PySpark DataFrame to a DMO in Data Cloud.
|
|
541
|
+
|
|
542
|
+
Args:
|
|
543
|
+
name: The name of the DMO to write to.
|
|
544
|
+
dataframe: The PySpark DataFrame to write.
|
|
545
|
+
write_mode: The write mode to use for writing to the DMO.
|
|
546
|
+
"""
|
|
547
|
+
self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DLO)
|
|
548
|
+
return self._writer.write_to_dmo(name, dataframe, write_mode, **kwargs) # type: ignore[no-any-return]
|
|
549
|
+
|
|
550
|
+
|
|
551
|
+
class StreamingClient(_BaseClient):
|
|
552
|
+
"""Entrypoint for streaming (``DELTA_SYNC``) access to Data Cloud objects.
|
|
553
|
+
|
|
554
|
+
This is the streaming counterpart to :class:`Client`. Instead of reading and
|
|
555
|
+
writing bounded snapshots, it reads a DLO/DMO change feed as a streaming
|
|
556
|
+
DataFrame and writes the transformed stream back via a ``StreamingQuery``.
|
|
557
|
+
"""
|
|
558
|
+
|
|
559
|
+
_instance: ClassVar[Optional[StreamingClient]] = None
|
|
560
|
+
|
|
561
|
+
def read_dlo_deltas(self) -> PySparkDataFrame:
|
|
562
|
+
"""Read the streaming change feed (deltas) for a DLO from Data Cloud.
|
|
563
|
+
|
|
564
|
+
For use in a streaming (``DELTA_SYNC``) BYOC transform. Returns a
|
|
565
|
+
streaming DataFrame whose rows carry the change-feed metadata columns
|
|
566
|
+
(``_record_type``, ``_commit_*``) alongside the source columns.
|
|
567
|
+
|
|
568
|
+
Returns:
|
|
569
|
+
A streaming PySpark DataFrame over the DLO change feed.
|
|
570
|
+
"""
|
|
571
|
+
self._record_dlo_access(_streaming_source_name())
|
|
572
|
+
return self._reader.read_dlo_deltas() # type: ignore[no-any-return]
|
|
573
|
+
|
|
574
|
+
def read_dmo_deltas(self) -> PySparkDataFrame:
|
|
575
|
+
"""Read the streaming change feed (deltas) for a DMO from Data Cloud.
|
|
576
|
+
|
|
577
|
+
Returns:
|
|
578
|
+
A streaming PySpark DataFrame over the DMO change feed.
|
|
579
|
+
"""
|
|
580
|
+
self._record_dmo_access(_streaming_source_name())
|
|
581
|
+
return self._reader.read_dmo_deltas() # type: ignore[no-any-return]
|
|
582
|
+
|
|
583
|
+
def write_dlo_deltas(
|
|
584
|
+
self, name: str, dataframe: PySparkDataFrame, **kwargs
|
|
585
|
+
) -> StreamingQuery:
|
|
586
|
+
"""Write a streaming DataFrame of deltas to a DLO in Data Cloud.
|
|
587
|
+
|
|
588
|
+
Starts a streaming query that writes each micro-batch to the
|
|
589
|
+
target DLO and returns the ``StreamingQuery`` handle; the caller
|
|
590
|
+
typically calls ``query.awaitTermination()``.
|
|
591
|
+
|
|
592
|
+
Args:
|
|
593
|
+
name: The name of the DLO to write to.
|
|
594
|
+
dataframe: The streaming PySpark DataFrame to write.
|
|
595
|
+
|
|
596
|
+
Returns:
|
|
597
|
+
The started ``StreamingQuery``.
|
|
598
|
+
"""
|
|
599
|
+
self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO)
|
|
600
|
+
return self._writer.write_dlo_deltas(name, dataframe, **kwargs) # type: ignore[no-any-return]
|
datacustomcode/config.py
CHANGED
|
@@ -89,6 +89,9 @@ class ClientConfig(BaseConfig):
|
|
|
89
89
|
spark_provider_config: Union[
|
|
90
90
|
SparkProviderConfig[BaseSparkSessionProvider], None
|
|
91
91
|
] = None
|
|
92
|
+
# Source object name for a streaming (DELTA_SYNC) transform, populated by
|
|
93
|
+
# ``run_entrypoint`` from config.json's ``streamingSource`` field
|
|
94
|
+
streaming_source: Union[str, None] = None
|
|
92
95
|
|
|
93
96
|
def update(self, other: ClientConfig) -> ClientConfig:
|
|
94
97
|
"""Merge this ClientConfig with another, respecting force flags.
|
|
@@ -116,6 +119,8 @@ class ClientConfig(BaseConfig):
|
|
|
116
119
|
self.spark_provider_config = merge(
|
|
117
120
|
self.spark_provider_config, other.spark_provider_config
|
|
118
121
|
)
|
|
122
|
+
if other.streaming_source is not None:
|
|
123
|
+
self.streaming_source = other.streaming_source
|
|
119
124
|
return self
|
|
120
125
|
|
|
121
126
|
|
datacustomcode/constants.py
CHANGED
|
@@ -38,6 +38,14 @@ USE_IN_FEATURE_MAPPING_FOR_CONNECT_API = {
|
|
|
38
38
|
"SearchIndexChunking": "UnstructuredChunking",
|
|
39
39
|
}
|
|
40
40
|
|
|
41
|
+
# Script (data transform) invoke options
|
|
42
|
+
SCRIPT_USE_IN_FEATURE_BATCH = "BatchTransform"
|
|
43
|
+
SCRIPT_USE_IN_FEATURE_STREAMING = "StreamingTransform"
|
|
44
|
+
SCRIPT_USE_IN_FEATURE_OPTIONS = [
|
|
45
|
+
SCRIPT_USE_IN_FEATURE_BATCH,
|
|
46
|
+
SCRIPT_USE_IN_FEATURE_STREAMING,
|
|
47
|
+
]
|
|
48
|
+
|
|
41
49
|
# Pydantic request/response type names to feature names
|
|
42
50
|
REQUEST_TYPE_TO_FEATURE = {
|
|
43
51
|
"SearchIndexChunkingV1Request": "SearchIndexChunking",
|
datacustomcode/deploy.py
CHANGED
|
@@ -14,6 +14,7 @@
|
|
|
14
14
|
# limitations under the License.
|
|
15
15
|
from __future__ import annotations
|
|
16
16
|
|
|
17
|
+
import copy
|
|
17
18
|
from html import unescape
|
|
18
19
|
import json
|
|
19
20
|
import os
|
|
@@ -41,6 +42,7 @@ from datacustomcode.scan import find_base_directory, get_package_type
|
|
|
41
42
|
|
|
42
43
|
DATA_CUSTOM_CODE_PATH = "services/data/v63.0/ssot/data-custom-code"
|
|
43
44
|
DATA_TRANSFORMS_PATH = "services/data/v63.0/ssot/data-transforms"
|
|
45
|
+
DATA_CUSTOM_CODE_INVOKE_OPTIONS_PATH = "services/data/v67.0/ssot/data-custom-code"
|
|
44
46
|
WAIT_FOR_DEPLOYMENT_TIMEOUT = 3000
|
|
45
47
|
|
|
46
48
|
# Available compute types for Data Cloud deployments.
|
|
@@ -108,6 +110,7 @@ class CodeExtensionMetadata(BaseModel):
|
|
|
108
110
|
computeType: str
|
|
109
111
|
codeType: str
|
|
110
112
|
functionInvokeOptions: Union[list[str], None] = None
|
|
113
|
+
invokeOptions: Union[list[str], None] = None
|
|
111
114
|
|
|
112
115
|
def __init__(self, **data):
|
|
113
116
|
name = data.get("name", "")
|
|
@@ -200,7 +203,14 @@ def create_deployment(
|
|
|
200
203
|
access_token: AccessTokenResponse, metadata: CodeExtensionMetadata
|
|
201
204
|
) -> CreateDeploymentResponse:
|
|
202
205
|
"""Create a custom code deployment in the DataCloud."""
|
|
203
|
-
|
|
206
|
+
# invokeOptions only binds at v67.0; route there when it is set so the
|
|
207
|
+
# option isn't silently dropped. Everything else stays on v63.0.
|
|
208
|
+
code_custom_code_path = (
|
|
209
|
+
DATA_CUSTOM_CODE_INVOKE_OPTIONS_PATH
|
|
210
|
+
if metadata.invokeOptions
|
|
211
|
+
else DATA_CUSTOM_CODE_PATH
|
|
212
|
+
)
|
|
213
|
+
url = _join_strip_url(access_token.instance_url, code_custom_code_path)
|
|
204
214
|
body = dict[str, Any](
|
|
205
215
|
{
|
|
206
216
|
"label": metadata.name,
|
|
@@ -213,6 +223,8 @@ def create_deployment(
|
|
|
213
223
|
)
|
|
214
224
|
if metadata.functionInvokeOptions:
|
|
215
225
|
body["functionInvokeOptions"] = metadata.functionInvokeOptions
|
|
226
|
+
if metadata.invokeOptions:
|
|
227
|
+
body["invokeOptions"] = metadata.invokeOptions
|
|
216
228
|
logger.debug(f"Creating deployment {metadata.name}...")
|
|
217
229
|
try:
|
|
218
230
|
response = _make_api_call(
|
|
@@ -388,6 +400,30 @@ class DataTransformConfig(BaseConfig):
|
|
|
388
400
|
dataspace: str
|
|
389
401
|
permissions: Permissions
|
|
390
402
|
dataObjects: Optional[list[DataObject]] = None
|
|
403
|
+
streamingSource: Optional[StreamingSource] = None
|
|
404
|
+
|
|
405
|
+
@property
|
|
406
|
+
def is_streaming(self) -> bool:
|
|
407
|
+
return self.streamingSource is not None
|
|
408
|
+
|
|
409
|
+
@model_validator(mode="after")
|
|
410
|
+
def _validate_layers(self) -> "DataTransformConfig":
|
|
411
|
+
read_is_dlo = isinstance(self.permissions.read, DloPermission)
|
|
412
|
+
write_is_dlo = isinstance(self.permissions.write, DloPermission)
|
|
413
|
+
if self.is_streaming:
|
|
414
|
+
if not write_is_dlo:
|
|
415
|
+
raise ValueError(
|
|
416
|
+
"A streaming transform must write to a DLO "
|
|
417
|
+
"(permissions.write must be a 'dlo' entry)."
|
|
418
|
+
)
|
|
419
|
+
elif read_is_dlo != write_is_dlo:
|
|
420
|
+
raise ValueError(
|
|
421
|
+
"permissions.read and permissions.write must both reference "
|
|
422
|
+
"DLOs or both reference DMOs (got "
|
|
423
|
+
f"read={type(self.permissions.read).__name__}, "
|
|
424
|
+
f"write={type(self.permissions.write).__name__})"
|
|
425
|
+
)
|
|
426
|
+
return self
|
|
391
427
|
|
|
392
428
|
|
|
393
429
|
class FunctionConfig(BaseConfig):
|
|
@@ -402,23 +438,15 @@ class DmoPermission(BaseModel):
|
|
|
402
438
|
dmo: list[str]
|
|
403
439
|
|
|
404
440
|
|
|
441
|
+
class StreamingSource(BaseModel):
|
|
442
|
+
type: str
|
|
443
|
+
name: str
|
|
444
|
+
|
|
445
|
+
|
|
405
446
|
class Permissions(BaseModel):
|
|
406
447
|
read: Union[DloPermission, DmoPermission]
|
|
407
448
|
write: Union[DloPermission, DmoPermission]
|
|
408
449
|
|
|
409
|
-
@model_validator(mode="after")
|
|
410
|
-
def _no_mixed_layers(self) -> "Permissions":
|
|
411
|
-
read_is_dlo = isinstance(self.read, DloPermission)
|
|
412
|
-
write_is_dlo = isinstance(self.write, DloPermission)
|
|
413
|
-
if read_is_dlo != write_is_dlo:
|
|
414
|
-
raise ValueError(
|
|
415
|
-
"permissions.read and permissions.write must both reference "
|
|
416
|
-
"DLOs or both reference DMOs (got "
|
|
417
|
-
f"read={type(self.read).__name__}, "
|
|
418
|
-
f"write={type(self.write).__name__})"
|
|
419
|
-
)
|
|
420
|
-
return self
|
|
421
|
-
|
|
422
450
|
|
|
423
451
|
def _permission_entries(perm: Union[DloPermission, DmoPermission]) -> list[str]:
|
|
424
452
|
"""Return the list of object names regardless of layer (DLO or DMO)."""
|
|
@@ -490,7 +518,9 @@ def create_data_transform(
|
|
|
490
518
|
) -> dict:
|
|
491
519
|
"""Create a data transform in the DataCloud."""
|
|
492
520
|
script_name = metadata.name
|
|
493
|
-
|
|
521
|
+
# Deep copy: the template's nested nodes/sources/macros dicts would
|
|
522
|
+
# otherwise be shared across calls and accumulate entries between deploys.
|
|
523
|
+
request_hydrated = copy.deepcopy(DATA_TRANSFORM_REQUEST_TEMPLATE)
|
|
494
524
|
|
|
495
525
|
# Add nodes for each write entry (DLO or DMO)
|
|
496
526
|
for i, name in enumerate(
|
|
@@ -533,7 +563,7 @@ def create_data_transform(
|
|
|
533
563
|
"definition": definition,
|
|
534
564
|
"label": f"{metadata.name}",
|
|
535
565
|
"name": f"{metadata.name}",
|
|
536
|
-
"type": "BATCH",
|
|
566
|
+
"type": "STREAMING" if data_transform_config.is_streaming else "BATCH",
|
|
537
567
|
"dataSpaceName": data_transform_config.dataspace,
|
|
538
568
|
}
|
|
539
569
|
|
|
@@ -616,9 +646,18 @@ def deploy_full(
|
|
|
616
646
|
callback=None,
|
|
617
647
|
) -> AccessTokenResponse:
|
|
618
648
|
"""Deploy a data transform in the DataCloud."""
|
|
649
|
+
from datacustomcode.constants import SCRIPT_USE_IN_FEATURE_STREAMING
|
|
650
|
+
|
|
619
651
|
# prepare payload
|
|
620
652
|
config = get_config(directory)
|
|
621
653
|
|
|
654
|
+
if (
|
|
655
|
+
isinstance(config, DataTransformConfig)
|
|
656
|
+
and config.is_streaming
|
|
657
|
+
and not metadata.invokeOptions
|
|
658
|
+
):
|
|
659
|
+
metadata.invokeOptions = [SCRIPT_USE_IN_FEATURE_STREAMING]
|
|
660
|
+
|
|
622
661
|
# create deployment and upload payload
|
|
623
662
|
deployment = create_deployment(access_token, metadata)
|
|
624
663
|
zip(directory, docker_network, metadata.codeType)
|