salesforce-data-customcode 6.1.0.dev2__py3-none-any.whl → 6.1.0.dev4__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 +212 -94
- datacustomcode/config.py +5 -0
- datacustomcode/constants.py +8 -0
- datacustomcode/deploy.py +55 -16
- datacustomcode/io/reader/base.py +42 -0
- datacustomcode/io/writer/base.py +33 -0
- datacustomcode/named_credential/__init__.py +2 -0
- datacustomcode/named_credential/direct/transport.py +6 -6
- datacustomcode/named_credential/errors.py +36 -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.1.0.dev2.dist-info → salesforce_data_customcode-6.1.0.dev4.dist-info}/METADATA +40 -3
- {salesforce_data_customcode-6.1.0.dev2.dist-info → salesforce_data_customcode-6.1.0.dev4.dist-info}/RECORD +20 -18
- {salesforce_data_customcode-6.1.0.dev2.dist-info → salesforce_data_customcode-6.1.0.dev4.dist-info}/WHEEL +0 -0
- {salesforce_data_customcode-6.1.0.dev2.dist-info → salesforce_data_customcode-6.1.0.dev4.dist-info}/entry_points.txt +0 -0
- {salesforce_data_customcode-6.1.0.dev2.dist-info → salesforce_data_customcode-6.1.0.dev4.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
|
|
@@ -35,7 +37,12 @@ from datacustomcode.spark.default import DefaultSparkSessionProvider
|
|
|
35
37
|
if TYPE_CHECKING:
|
|
36
38
|
from pathlib import Path
|
|
37
39
|
|
|
38
|
-
from pyspark.sql import
|
|
40
|
+
from pyspark.sql import (
|
|
41
|
+
Column,
|
|
42
|
+
DataFrame as PySparkDataFrame,
|
|
43
|
+
SparkSession,
|
|
44
|
+
)
|
|
45
|
+
from pyspark.sql.streaming import StreamingQuery
|
|
39
46
|
|
|
40
47
|
from datacustomcode.einstein_predictions.spark_base import SparkEinsteinPredictions
|
|
41
48
|
from datacustomcode.einstein_predictions.types import PredictionType
|
|
@@ -48,6 +55,40 @@ if TYPE_CHECKING:
|
|
|
48
55
|
from datacustomcode.spark.base import BaseSparkSessionProvider
|
|
49
56
|
|
|
50
57
|
|
|
58
|
+
def _streaming_source_name() -> str:
|
|
59
|
+
"""Return the streaming transform's read-source name.
|
|
60
|
+
|
|
61
|
+
Resolved from ``config.streaming_source``, which ``run_entrypoint``
|
|
62
|
+
populates from config.json's ``streamingSource`` field.
|
|
63
|
+
|
|
64
|
+
Raises:
|
|
65
|
+
RuntimeError: If no ``streaming_source`` has been configured (e.g. the
|
|
66
|
+
transform's config.json has no ``streamingSource`` field).
|
|
67
|
+
"""
|
|
68
|
+
source = config.streaming_source
|
|
69
|
+
if not source:
|
|
70
|
+
raise RuntimeError(
|
|
71
|
+
"No streaming source configured. A streaming transform must declare "
|
|
72
|
+
"its read source in config.json under 'streamingSource'."
|
|
73
|
+
)
|
|
74
|
+
return source
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
def _active_client() -> "_BaseClient":
|
|
78
|
+
"""Return the client backing the module-level Spark column helpers.
|
|
79
|
+
|
|
80
|
+
Prefers an already-initialized singleton so a streaming job reuses its
|
|
81
|
+
:class:`StreamingClient` (and a batch job its :class:`Client`) rather than
|
|
82
|
+
forcing an unrelated client into existence. Falls back to building the
|
|
83
|
+
batch :class:`Client` when neither has been created yet.
|
|
84
|
+
"""
|
|
85
|
+
if Client._instance is not None:
|
|
86
|
+
return Client._instance
|
|
87
|
+
if StreamingClient._instance is not None:
|
|
88
|
+
return StreamingClient._instance
|
|
89
|
+
return Client()
|
|
90
|
+
|
|
91
|
+
|
|
51
92
|
def _build_spark_llm_gateway() -> "SparkLLMGateway":
|
|
52
93
|
"""Instantiate the SDK-configured :class:`SparkLLMGateway`.
|
|
53
94
|
|
|
@@ -103,7 +144,7 @@ def llm_gateway_generate_text_col(
|
|
|
103
144
|
the generated text; on failure, ``status == "ERROR"`` and the
|
|
104
145
|
``error_*`` fields carry diagnostic detail.
|
|
105
146
|
"""
|
|
106
|
-
gateway =
|
|
147
|
+
gateway = _active_client()._get_spark_llm_gateway()
|
|
107
148
|
return gateway.llm_gateway_generate_text_col(template, values, model_id=model_id)
|
|
108
149
|
|
|
109
150
|
|
|
@@ -180,7 +221,7 @@ def einstein_predict_col(
|
|
|
180
221
|
the JSON-serialized prediction payload; on failure, ``status ==
|
|
181
222
|
"ERROR"`` and the ``error_*`` fields carry diagnostic detail.
|
|
182
223
|
"""
|
|
183
|
-
predictions =
|
|
224
|
+
predictions = _active_client()._get_spark_einstein_predictions()
|
|
184
225
|
return predictions.einstein_predict_col(
|
|
185
226
|
model_api_name, prediction_type, features, settings=settings
|
|
186
227
|
)
|
|
@@ -258,26 +299,22 @@ class DataCloudAccessLayerException(Exception):
|
|
|
258
299
|
return msg
|
|
259
300
|
|
|
260
301
|
|
|
261
|
-
|
|
262
|
-
|
|
302
|
+
_ClientT = TypeVar("_ClientT", bound="_BaseClient")
|
|
303
|
+
|
|
304
|
+
|
|
305
|
+
class _BaseClient:
|
|
306
|
+
"""Shared machinery for the Data Cloud client singletons.
|
|
263
307
|
|
|
264
|
-
|
|
265
|
-
|
|
266
|
-
can read from DLOs and write to DLOs, read from DMOs and write to DMOs, but you
|
|
267
|
-
cannot read from DLOs and write to DMOs or read from DMOs and write to DLOs.
|
|
268
|
-
Furthermore you cannot mix during merging tables. This class is a singleton to
|
|
269
|
-
prevent accidental mixing of DLOs and DMOs.
|
|
308
|
+
Holds the wiring common to :class:`Client` (batch) and
|
|
309
|
+
:class:`StreamingClient`
|
|
270
310
|
|
|
271
|
-
|
|
272
|
-
|
|
273
|
-
behavior once deployed to Data Cloud. By default, the client intercepts all
|
|
274
|
-
read/write operations and mocks access to Data Cloud. For example, during
|
|
275
|
-
writing, we print to the console instead of writing to Data Cloud.
|
|
311
|
+
This base class is not meant to be instantiated directly; use
|
|
312
|
+
:class:`Client` or :class:`StreamingClient`.
|
|
276
313
|
|
|
277
314
|
Args:
|
|
278
|
-
finder: Find a file path
|
|
279
315
|
reader: A custom reader to use for reading Data Cloud objects.
|
|
280
316
|
writer: A custom writer to use for writing Data Cloud objects.
|
|
317
|
+
spark_provider: Optional custom :class:`BaseSparkSessionProvider`.
|
|
281
318
|
spark_llm_gateway: Optional custom :class:`SparkLLMGateway`.
|
|
282
319
|
spark_einstein_predictions: Optional custom
|
|
283
320
|
:class:`SparkEinsteinPredictions`.
|
|
@@ -291,7 +328,18 @@ class Client:
|
|
|
291
328
|
>>> answer = client.llm_gateway_generate_text("Generate a greeting message")
|
|
292
329
|
"""
|
|
293
330
|
|
|
294
|
-
_instance:
|
|
331
|
+
# Each concrete subclass gets its own ``_instance`` slot: reads fall through
|
|
332
|
+
# to this base default of ``None``, but ``cls._instance = ...`` in __new__
|
|
333
|
+
# always writes to the subclass, so ``Client`` and ``StreamingClient`` never
|
|
334
|
+
# share an instance.
|
|
335
|
+
_instance: ClassVar[Optional[_BaseClient]] = None
|
|
336
|
+
# Process-wide Spark session shared across BOTH client types. Unlike
|
|
337
|
+
# ``_instance``, this is written via ``_BaseClient._shared_spark`` (never
|
|
338
|
+
# ``cls._shared_spark``), so the slot lives on the base class and a
|
|
339
|
+
# ``Client`` and a ``StreamingClient`` in the same process reuse one session
|
|
340
|
+
# — and therefore one underlying connection — instead of opening two
|
|
341
|
+
# containing differing state
|
|
342
|
+
_shared_spark: ClassVar[Optional[SparkSession]] = None
|
|
295
343
|
_reader: BaseDataCloudReader
|
|
296
344
|
_writer: BaseDataCloudWriter
|
|
297
345
|
_file: DefaultFindFilePath
|
|
@@ -302,7 +350,7 @@ class Client:
|
|
|
302
350
|
_code_type: str
|
|
303
351
|
|
|
304
352
|
def __new__(
|
|
305
|
-
cls,
|
|
353
|
+
cls: type[_ClientT],
|
|
306
354
|
reader: Optional[BaseDataCloudReader] = None,
|
|
307
355
|
writer: Optional[BaseDataCloudWriter] = None,
|
|
308
356
|
spark_provider: Optional[BaseSparkSessionProvider] = None,
|
|
@@ -310,31 +358,38 @@ class Client:
|
|
|
310
358
|
spark_einstein_predictions: Optional[SparkEinsteinPredictions] = None,
|
|
311
359
|
spark_named_credential: Optional[SparkNamedCredential] = None,
|
|
312
360
|
code_type: str = "script",
|
|
313
|
-
) ->
|
|
361
|
+
) -> _ClientT:
|
|
314
362
|
|
|
315
363
|
if cls._instance is None:
|
|
316
|
-
|
|
317
|
-
|
|
318
|
-
|
|
319
|
-
|
|
364
|
+
instance = super().__new__(cls)
|
|
365
|
+
instance._spark_llm_gateway = spark_llm_gateway
|
|
366
|
+
instance._spark_einstein_predictions = spark_einstein_predictions
|
|
367
|
+
instance._spark_named_credential = spark_named_credential
|
|
320
368
|
# Initialize Readers and Writers from config
|
|
321
369
|
# and/or provided reader and writer
|
|
322
370
|
if reader is None or writer is None:
|
|
323
|
-
# We need a spark because we will initialize readers and writers
|
|
324
|
-
if
|
|
325
|
-
|
|
326
|
-
|
|
327
|
-
|
|
328
|
-
|
|
329
|
-
provider: BaseSparkSessionProvider
|
|
330
|
-
if spark_provider is not None:
|
|
331
|
-
provider = spark_provider
|
|
332
|
-
elif config.spark_provider_config is not None:
|
|
333
|
-
provider = config.spark_provider_config.to_object()
|
|
371
|
+
# We need a spark because we will initialize readers and writers.
|
|
372
|
+
# Reuse the process-wide session if one client already built it,
|
|
373
|
+
# so a Client and a StreamingClient share a single connection.
|
|
374
|
+
if _BaseClient._shared_spark is not None:
|
|
375
|
+
spark = _BaseClient._shared_spark
|
|
334
376
|
else:
|
|
335
|
-
|
|
336
|
-
|
|
337
|
-
|
|
377
|
+
if config.spark_config is None:
|
|
378
|
+
raise ValueError(
|
|
379
|
+
"Spark config is required when reader/writer is not "
|
|
380
|
+
"provided"
|
|
381
|
+
)
|
|
382
|
+
|
|
383
|
+
provider: BaseSparkSessionProvider
|
|
384
|
+
if spark_provider is not None:
|
|
385
|
+
provider = spark_provider
|
|
386
|
+
elif config.spark_provider_config is not None:
|
|
387
|
+
provider = config.spark_provider_config.to_object()
|
|
388
|
+
else:
|
|
389
|
+
provider = DefaultSparkSessionProvider()
|
|
390
|
+
|
|
391
|
+
spark = provider.get_session(config.spark_config)
|
|
392
|
+
_BaseClient._shared_spark = spark
|
|
338
393
|
|
|
339
394
|
if config.reader_config is None and reader is None:
|
|
340
395
|
raise ValueError(
|
|
@@ -357,66 +412,17 @@ class Client:
|
|
|
357
412
|
else:
|
|
358
413
|
writer_init = writer
|
|
359
414
|
|
|
360
|
-
|
|
361
|
-
|
|
362
|
-
|
|
363
|
-
|
|
415
|
+
instance._reader = reader_init
|
|
416
|
+
instance._writer = writer_init
|
|
417
|
+
instance._file = DefaultFindFilePath()
|
|
418
|
+
instance._data_layer_history = {
|
|
364
419
|
DataCloudObjectType.DLO: set(),
|
|
365
420
|
DataCloudObjectType.DMO: set(),
|
|
366
421
|
}
|
|
367
|
-
|
|
422
|
+
cls._instance = instance
|
|
423
|
+
elif reader is not None or writer is not None:
|
|
368
424
|
raise ValueError("Cannot set reader or writer after client is initialized")
|
|
369
|
-
return cls._instance
|
|
370
|
-
|
|
371
|
-
def read_dlo(self, name: str) -> PySparkDataFrame:
|
|
372
|
-
"""Read a DLO from Data Cloud.
|
|
373
|
-
|
|
374
|
-
Args:
|
|
375
|
-
name: The name of the DLO to read.
|
|
376
|
-
|
|
377
|
-
Returns:
|
|
378
|
-
A PySpark DataFrame containing the DLO data.
|
|
379
|
-
"""
|
|
380
|
-
self._record_dlo_access(name)
|
|
381
|
-
return self._reader.read_dlo(name) # type: ignore[no-any-return]
|
|
382
|
-
|
|
383
|
-
def read_dmo(self, name: str) -> PySparkDataFrame:
|
|
384
|
-
"""Read a DMO from Data Cloud.
|
|
385
|
-
|
|
386
|
-
Args:
|
|
387
|
-
name: The name of the DMO to read.
|
|
388
|
-
|
|
389
|
-
Returns:
|
|
390
|
-
A PySpark DataFrame containing the DMO data.
|
|
391
|
-
"""
|
|
392
|
-
self._record_dmo_access(name)
|
|
393
|
-
return self._reader.read_dmo(name) # type: ignore[no-any-return]
|
|
394
|
-
|
|
395
|
-
def write_to_dlo(
|
|
396
|
-
self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs
|
|
397
|
-
) -> None:
|
|
398
|
-
"""Write a PySpark DataFrame to a DLO in Data Cloud.
|
|
399
|
-
|
|
400
|
-
Args:
|
|
401
|
-
name: The name of the DLO to write to.
|
|
402
|
-
dataframe: The PySpark DataFrame to write.
|
|
403
|
-
write_mode: The write mode to use for writing to the DLO.
|
|
404
|
-
"""
|
|
405
|
-
self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO)
|
|
406
|
-
return self._writer.write_to_dlo(name, dataframe, write_mode, **kwargs) # type: ignore[no-any-return]
|
|
407
|
-
|
|
408
|
-
def write_to_dmo(
|
|
409
|
-
self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs
|
|
410
|
-
) -> None:
|
|
411
|
-
"""Write a PySpark DataFrame to a DMO in Data Cloud.
|
|
412
|
-
|
|
413
|
-
Args:
|
|
414
|
-
name: The name of the DMO to write to.
|
|
415
|
-
dataframe: The PySpark DataFrame to write.
|
|
416
|
-
write_mode: The write mode to use for writing to the DMO.
|
|
417
|
-
"""
|
|
418
|
-
self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DLO)
|
|
419
|
-
return self._writer.write_to_dmo(name, dataframe, write_mode, **kwargs) # type: ignore[no-any-return]
|
|
425
|
+
return cast(_ClientT, cls._instance)
|
|
420
426
|
|
|
421
427
|
def find_file_path(self, file_name: str) -> Path:
|
|
422
428
|
"""Resolve a bundled file shipped in the package to an absolute path.
|
|
@@ -579,3 +585,115 @@ class Client:
|
|
|
579
585
|
|
|
580
586
|
def _record_dmo_access(self, name: str) -> None:
|
|
581
587
|
self._data_layer_history[DataCloudObjectType.DMO].add(name)
|
|
588
|
+
|
|
589
|
+
|
|
590
|
+
class Client(_BaseClient):
|
|
591
|
+
"""Entrypoint for batch access to Data Cloud objects.
|
|
592
|
+
|
|
593
|
+
This is the object used to read and write bounded snapshots of Data Cloud
|
|
594
|
+
DLOs and DMOs.
|
|
595
|
+
"""
|
|
596
|
+
|
|
597
|
+
_instance: ClassVar[Optional[Client]] = None
|
|
598
|
+
|
|
599
|
+
def read_dlo(self, name: str) -> PySparkDataFrame:
|
|
600
|
+
"""Read a DLO from Data Cloud.
|
|
601
|
+
|
|
602
|
+
Args:
|
|
603
|
+
name: The name of the DLO to read.
|
|
604
|
+
|
|
605
|
+
Returns:
|
|
606
|
+
A PySpark DataFrame containing the DLO data.
|
|
607
|
+
"""
|
|
608
|
+
self._record_dlo_access(name)
|
|
609
|
+
return self._reader.read_dlo(name) # type: ignore[no-any-return]
|
|
610
|
+
|
|
611
|
+
def read_dmo(self, name: str) -> PySparkDataFrame:
|
|
612
|
+
"""Read a DMO from Data Cloud.
|
|
613
|
+
|
|
614
|
+
Args:
|
|
615
|
+
name: The name of the DMO to read.
|
|
616
|
+
|
|
617
|
+
Returns:
|
|
618
|
+
A PySpark DataFrame containing the DMO data.
|
|
619
|
+
"""
|
|
620
|
+
self._record_dmo_access(name)
|
|
621
|
+
return self._reader.read_dmo(name) # type: ignore[no-any-return]
|
|
622
|
+
|
|
623
|
+
def write_to_dlo(
|
|
624
|
+
self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs
|
|
625
|
+
) -> None:
|
|
626
|
+
"""Write a PySpark DataFrame to a DLO in Data Cloud.
|
|
627
|
+
|
|
628
|
+
Args:
|
|
629
|
+
name: The name of the DLO to write to.
|
|
630
|
+
dataframe: The PySpark DataFrame to write.
|
|
631
|
+
write_mode: The write mode to use for writing to the DLO.
|
|
632
|
+
"""
|
|
633
|
+
self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO)
|
|
634
|
+
return self._writer.write_to_dlo(name, dataframe, write_mode, **kwargs) # type: ignore[no-any-return]
|
|
635
|
+
|
|
636
|
+
def write_to_dmo(
|
|
637
|
+
self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs
|
|
638
|
+
) -> None:
|
|
639
|
+
"""Write a PySpark DataFrame to a DMO in Data Cloud.
|
|
640
|
+
|
|
641
|
+
Args:
|
|
642
|
+
name: The name of the DMO to write to.
|
|
643
|
+
dataframe: The PySpark DataFrame to write.
|
|
644
|
+
write_mode: The write mode to use for writing to the DMO.
|
|
645
|
+
"""
|
|
646
|
+
self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DLO)
|
|
647
|
+
return self._writer.write_to_dmo(name, dataframe, write_mode, **kwargs) # type: ignore[no-any-return]
|
|
648
|
+
|
|
649
|
+
|
|
650
|
+
class StreamingClient(_BaseClient):
|
|
651
|
+
"""Entrypoint for streaming (``DELTA_SYNC``) access to Data Cloud objects.
|
|
652
|
+
|
|
653
|
+
This is the streaming counterpart to :class:`Client`. Instead of reading and
|
|
654
|
+
writing bounded snapshots, it reads a DLO/DMO change feed as a streaming
|
|
655
|
+
DataFrame and writes the transformed stream back via a ``StreamingQuery``.
|
|
656
|
+
"""
|
|
657
|
+
|
|
658
|
+
_instance: ClassVar[Optional[StreamingClient]] = None
|
|
659
|
+
|
|
660
|
+
def read_dlo_deltas(self) -> PySparkDataFrame:
|
|
661
|
+
"""Read the streaming change feed (deltas) for a DLO from Data Cloud.
|
|
662
|
+
|
|
663
|
+
For use in a streaming (``DELTA_SYNC``) BYOC transform. Returns a
|
|
664
|
+
streaming DataFrame whose rows carry the change-feed metadata columns
|
|
665
|
+
(``_record_type``, ``_commit_*``) alongside the source columns.
|
|
666
|
+
|
|
667
|
+
Returns:
|
|
668
|
+
A streaming PySpark DataFrame over the DLO change feed.
|
|
669
|
+
"""
|
|
670
|
+
self._record_dlo_access(_streaming_source_name())
|
|
671
|
+
return self._reader.read_dlo_deltas() # type: ignore[no-any-return]
|
|
672
|
+
|
|
673
|
+
def read_dmo_deltas(self) -> PySparkDataFrame:
|
|
674
|
+
"""Read the streaming change feed (deltas) for a DMO from Data Cloud.
|
|
675
|
+
|
|
676
|
+
Returns:
|
|
677
|
+
A streaming PySpark DataFrame over the DMO change feed.
|
|
678
|
+
"""
|
|
679
|
+
self._record_dmo_access(_streaming_source_name())
|
|
680
|
+
return self._reader.read_dmo_deltas() # type: ignore[no-any-return]
|
|
681
|
+
|
|
682
|
+
def write_dlo_deltas(
|
|
683
|
+
self, name: str, dataframe: PySparkDataFrame, **kwargs
|
|
684
|
+
) -> StreamingQuery:
|
|
685
|
+
"""Write a streaming DataFrame of deltas to a DLO in Data Cloud.
|
|
686
|
+
|
|
687
|
+
Starts a streaming query that writes each micro-batch to the
|
|
688
|
+
target DLO and returns the ``StreamingQuery`` handle; the caller
|
|
689
|
+
typically calls ``query.awaitTermination()``.
|
|
690
|
+
|
|
691
|
+
Args:
|
|
692
|
+
name: The name of the DLO to write to.
|
|
693
|
+
dataframe: The streaming PySpark DataFrame to write.
|
|
694
|
+
|
|
695
|
+
Returns:
|
|
696
|
+
The started ``StreamingQuery``.
|
|
697
|
+
"""
|
|
698
|
+
self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO)
|
|
699
|
+
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
|
|
@@ -44,6 +45,7 @@ from datacustomcode.scan import find_base_directory, get_package_type
|
|
|
44
45
|
|
|
45
46
|
DATA_CUSTOM_CODE_PATH = "services/data/v63.0/ssot/data-custom-code"
|
|
46
47
|
DATA_TRANSFORMS_PATH = "services/data/v63.0/ssot/data-transforms"
|
|
48
|
+
DATA_CUSTOM_CODE_INVOKE_OPTIONS_PATH = "services/data/v67.0/ssot/data-custom-code"
|
|
47
49
|
WAIT_FOR_DEPLOYMENT_TIMEOUT = 3000
|
|
48
50
|
|
|
49
51
|
# Available compute types for Data Cloud deployments.
|
|
@@ -111,6 +113,7 @@ class CodeExtensionMetadata(BaseModel):
|
|
|
111
113
|
computeType: str
|
|
112
114
|
codeType: str
|
|
113
115
|
functionInvokeOptions: Union[list[str], None] = None
|
|
116
|
+
invokeOptions: Union[list[str], None] = None
|
|
114
117
|
|
|
115
118
|
def __init__(self, **data):
|
|
116
119
|
name = data.get("name", "")
|
|
@@ -203,7 +206,14 @@ def create_deployment(
|
|
|
203
206
|
access_token: AccessTokenResponse, metadata: CodeExtensionMetadata
|
|
204
207
|
) -> CreateDeploymentResponse:
|
|
205
208
|
"""Create a custom code deployment in the DataCloud."""
|
|
206
|
-
|
|
209
|
+
# invokeOptions only binds at v67.0; route there when it is set so the
|
|
210
|
+
# option isn't silently dropped. Everything else stays on v63.0.
|
|
211
|
+
code_custom_code_path = (
|
|
212
|
+
DATA_CUSTOM_CODE_INVOKE_OPTIONS_PATH
|
|
213
|
+
if metadata.invokeOptions
|
|
214
|
+
else DATA_CUSTOM_CODE_PATH
|
|
215
|
+
)
|
|
216
|
+
url = _join_strip_url(access_token.instance_url, code_custom_code_path)
|
|
207
217
|
body = dict[str, Any](
|
|
208
218
|
{
|
|
209
219
|
"label": metadata.name,
|
|
@@ -216,6 +226,8 @@ def create_deployment(
|
|
|
216
226
|
)
|
|
217
227
|
if metadata.functionInvokeOptions:
|
|
218
228
|
body["functionInvokeOptions"] = metadata.functionInvokeOptions
|
|
229
|
+
if metadata.invokeOptions:
|
|
230
|
+
body["invokeOptions"] = metadata.invokeOptions
|
|
219
231
|
logger.debug(f"Creating deployment {metadata.name}...")
|
|
220
232
|
try:
|
|
221
233
|
response = _make_api_call(
|
|
@@ -393,6 +405,30 @@ class DataTransformConfig(BaseConfig):
|
|
|
393
405
|
dataspace: str
|
|
394
406
|
permissions: Permissions
|
|
395
407
|
dataObjects: Optional[list[DataObject]] = None
|
|
408
|
+
streamingSource: Optional[StreamingSource] = None
|
|
409
|
+
|
|
410
|
+
@property
|
|
411
|
+
def is_streaming(self) -> bool:
|
|
412
|
+
return self.streamingSource is not None
|
|
413
|
+
|
|
414
|
+
@model_validator(mode="after")
|
|
415
|
+
def _validate_layers(self) -> "DataTransformConfig":
|
|
416
|
+
read_is_dlo = isinstance(self.permissions.read, DloPermission)
|
|
417
|
+
write_is_dlo = isinstance(self.permissions.write, DloPermission)
|
|
418
|
+
if self.is_streaming:
|
|
419
|
+
if not write_is_dlo:
|
|
420
|
+
raise ValueError(
|
|
421
|
+
"A streaming transform must write to a DLO "
|
|
422
|
+
"(permissions.write must be a 'dlo' entry)."
|
|
423
|
+
)
|
|
424
|
+
elif read_is_dlo != write_is_dlo:
|
|
425
|
+
raise ValueError(
|
|
426
|
+
"permissions.read and permissions.write must both reference "
|
|
427
|
+
"DLOs or both reference DMOs (got "
|
|
428
|
+
f"read={type(self.permissions.read).__name__}, "
|
|
429
|
+
f"write={type(self.permissions.write).__name__})"
|
|
430
|
+
)
|
|
431
|
+
return self
|
|
396
432
|
|
|
397
433
|
|
|
398
434
|
class FunctionConfig(BaseConfig):
|
|
@@ -407,23 +443,15 @@ class DmoPermission(BaseModel):
|
|
|
407
443
|
dmo: list[str]
|
|
408
444
|
|
|
409
445
|
|
|
446
|
+
class StreamingSource(BaseModel):
|
|
447
|
+
type: str
|
|
448
|
+
name: str
|
|
449
|
+
|
|
450
|
+
|
|
410
451
|
class Permissions(BaseModel):
|
|
411
452
|
read: Union[DloPermission, DmoPermission]
|
|
412
453
|
write: Union[DloPermission, DmoPermission]
|
|
413
454
|
|
|
414
|
-
@model_validator(mode="after")
|
|
415
|
-
def _no_mixed_layers(self) -> "Permissions":
|
|
416
|
-
read_is_dlo = isinstance(self.read, DloPermission)
|
|
417
|
-
write_is_dlo = isinstance(self.write, DloPermission)
|
|
418
|
-
if read_is_dlo != write_is_dlo:
|
|
419
|
-
raise ValueError(
|
|
420
|
-
"permissions.read and permissions.write must both reference "
|
|
421
|
-
"DLOs or both reference DMOs (got "
|
|
422
|
-
f"read={type(self.read).__name__}, "
|
|
423
|
-
f"write={type(self.write).__name__})"
|
|
424
|
-
)
|
|
425
|
-
return self
|
|
426
|
-
|
|
427
455
|
|
|
428
456
|
def _permission_entries(perm: Union[DloPermission, DmoPermission]) -> list[str]:
|
|
429
457
|
"""Return the list of object names regardless of layer (DLO or DMO)."""
|
|
@@ -495,7 +523,9 @@ def create_data_transform(
|
|
|
495
523
|
) -> dict:
|
|
496
524
|
"""Create a data transform in the DataCloud."""
|
|
497
525
|
script_name = metadata.name
|
|
498
|
-
|
|
526
|
+
# Deep copy: the template's nested nodes/sources/macros dicts would
|
|
527
|
+
# otherwise be shared across calls and accumulate entries between deploys.
|
|
528
|
+
request_hydrated = copy.deepcopy(DATA_TRANSFORM_REQUEST_TEMPLATE)
|
|
499
529
|
|
|
500
530
|
# Add nodes for each write entry (DLO or DMO)
|
|
501
531
|
for i, name in enumerate(
|
|
@@ -538,7 +568,7 @@ def create_data_transform(
|
|
|
538
568
|
"definition": definition,
|
|
539
569
|
"label": f"{metadata.name}",
|
|
540
570
|
"name": f"{metadata.name}",
|
|
541
|
-
"type": "BATCH",
|
|
571
|
+
"type": "STREAMING" if data_transform_config.is_streaming else "BATCH",
|
|
542
572
|
"dataSpaceName": data_transform_config.dataspace,
|
|
543
573
|
}
|
|
544
574
|
|
|
@@ -621,9 +651,18 @@ def deploy_full(
|
|
|
621
651
|
callback=None,
|
|
622
652
|
) -> AccessTokenResponse:
|
|
623
653
|
"""Deploy a data transform in the DataCloud."""
|
|
654
|
+
from datacustomcode.constants import SCRIPT_USE_IN_FEATURE_STREAMING
|
|
655
|
+
|
|
624
656
|
# prepare payload
|
|
625
657
|
config = get_config(directory)
|
|
626
658
|
|
|
659
|
+
if (
|
|
660
|
+
isinstance(config, DataTransformConfig)
|
|
661
|
+
and config.is_streaming
|
|
662
|
+
and not metadata.invokeOptions
|
|
663
|
+
):
|
|
664
|
+
metadata.invokeOptions = [SCRIPT_USE_IN_FEATURE_STREAMING]
|
|
665
|
+
|
|
627
666
|
# create deployment and upload payload
|
|
628
667
|
deployment = create_deployment(access_token, metadata)
|
|
629
668
|
zip(directory, docker_network, metadata.codeType)
|