From 4c00881c6e0b3d9b87d389988abf1508fc6cbe41 Mon Sep 17 00:00:00 2001 From: Joshua Catt Date: Wed, 22 Jul 2026 15:02:18 -0400 Subject: [PATCH 1/2] Support creating and deploying streaming script code packages --- src/datacustomcode/cli.py | 33 +++- src/datacustomcode/constants.py | 8 + src/datacustomcode/deploy.py | 71 ++++++-- src/datacustomcode/scan.py | 196 ++++++++++++++++++---- src/datacustomcode/template.py | 14 +- tests/test_cli.py | 97 ++++++++++- tests/test_deploy.py | 286 ++++++++++++++++++++++++++++++++ tests/test_scan.py | 207 +++++++++++++++++++++++ 8 files changed, 860 insertions(+), 52 deletions(-) diff --git a/src/datacustomcode/cli.py b/src/datacustomcode/cli.py index 7e2cd00..3a03e6d 100644 --- a/src/datacustomcode/cli.py +++ b/src/datacustomcode/cli.py @@ -283,10 +283,19 @@ def deploy( ) @click.option( "--use-in-feature", - default="SearchIndexChunking", - help="Feature where this function will be used (only applicable for function).", + "-u", + default=None, + help=( + "Invoke option for this package. For scripts: 'BatchTransform' " + "(default) or 'StreamingTransform'. For functions: 'SearchIndexChunking'." + ), ) def init(directory: str, code_type: str, use_in_feature: Optional[str]): + from datacustomcode.constants import ( + SCRIPT_USE_IN_FEATURE_BATCH, + SCRIPT_USE_IN_FEATURE_OPTIONS, + SCRIPT_USE_IN_FEATURE_STREAMING, + ) from datacustomcode.scan import ( dc_config_json_from_file, update_config, @@ -294,9 +303,23 @@ def init(directory: str, code_type: str, use_in_feature: Optional[str]): ) from datacustomcode.template import copy_function_template, copy_script_template + streaming = False + if code_type == "script": + use_in_feature = use_in_feature or SCRIPT_USE_IN_FEATURE_BATCH + if use_in_feature not in SCRIPT_USE_IN_FEATURE_OPTIONS: + click.secho( + f"Error: Invalid --use-in-feature '{use_in_feature}' for a " + f"script. Valid options: {', '.join(SCRIPT_USE_IN_FEATURE_OPTIONS)}.", + fg="red", + ) + raise click.Abort() + streaming = use_in_feature == SCRIPT_USE_IN_FEATURE_STREAMING + else: + use_in_feature = use_in_feature or "SearchIndexChunking" + click.echo("Copying template to " + click.style(directory, fg="blue", bold=True)) if code_type == "script": - copy_script_template(directory) + copy_script_template(directory, streaming=streaming) elif code_type == "function": copy_function_template(directory, use_in_feature) 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]): sdk_config = {"type": code_type} write_sdk_config(directory, sdk_config) - config_json = dc_config_json_from_file(entrypoint_path, code_type) + config_json = dc_config_json_from_file( + entrypoint_path, code_type, streaming=streaming + ) with open(config_location, "w") as f: json.dump(config_json, f, indent=2) diff --git a/src/datacustomcode/constants.py b/src/datacustomcode/constants.py index 76b6a7c..fe29a5c 100644 --- a/src/datacustomcode/constants.py +++ b/src/datacustomcode/constants.py @@ -38,6 +38,14 @@ "SearchIndexChunking": "UnstructuredChunking", } +# Script (data transform) invoke options +SCRIPT_USE_IN_FEATURE_BATCH = "BatchTransform" +SCRIPT_USE_IN_FEATURE_STREAMING = "StreamingTransform" +SCRIPT_USE_IN_FEATURE_OPTIONS = [ + SCRIPT_USE_IN_FEATURE_BATCH, + SCRIPT_USE_IN_FEATURE_STREAMING, +] + # Pydantic request/response type names to feature names REQUEST_TYPE_TO_FEATURE = { "SearchIndexChunkingV1Request": "SearchIndexChunking", diff --git a/src/datacustomcode/deploy.py b/src/datacustomcode/deploy.py index e8c4ec4..e8555ba 100644 --- a/src/datacustomcode/deploy.py +++ b/src/datacustomcode/deploy.py @@ -14,6 +14,7 @@ # limitations under the License. from __future__ import annotations +import copy from html import unescape import json import os @@ -41,6 +42,7 @@ DATA_CUSTOM_CODE_PATH = "services/data/v63.0/ssot/data-custom-code" DATA_TRANSFORMS_PATH = "services/data/v63.0/ssot/data-transforms" +DATA_CUSTOM_CODE_INVOKE_OPTIONS_PATH = "services/data/v67.0/ssot/data-custom-code" WAIT_FOR_DEPLOYMENT_TIMEOUT = 3000 # Available compute types for Data Cloud deployments. @@ -108,6 +110,7 @@ class CodeExtensionMetadata(BaseModel): computeType: str codeType: str functionInvokeOptions: Union[list[str], None] = None + invokeOptions: Union[list[str], None] = None def __init__(self, **data): name = data.get("name", "") @@ -200,7 +203,14 @@ def create_deployment( access_token: AccessTokenResponse, metadata: CodeExtensionMetadata ) -> CreateDeploymentResponse: """Create a custom code deployment in the DataCloud.""" - url = _join_strip_url(access_token.instance_url, DATA_CUSTOM_CODE_PATH) + # invokeOptions only binds at v67.0; route there when it is set so the + # option isn't silently dropped. Everything else stays on v63.0. + code_custom_code_path = ( + DATA_CUSTOM_CODE_INVOKE_OPTIONS_PATH + if metadata.invokeOptions + else DATA_CUSTOM_CODE_PATH + ) + url = _join_strip_url(access_token.instance_url, code_custom_code_path) body = dict[str, Any]( { "label": metadata.name, @@ -213,6 +223,8 @@ def create_deployment( ) if metadata.functionInvokeOptions: body["functionInvokeOptions"] = metadata.functionInvokeOptions + if metadata.invokeOptions: + body["invokeOptions"] = metadata.invokeOptions logger.debug(f"Creating deployment {metadata.name}...") try: response = _make_api_call( @@ -388,6 +400,30 @@ class DataTransformConfig(BaseConfig): dataspace: str permissions: Permissions dataObjects: Optional[list[DataObject]] = None + streamingSource: Optional[StreamingSource] = None + + @property + def is_streaming(self) -> bool: + return self.streamingSource is not None + + @model_validator(mode="after") + def _validate_layers(self) -> "DataTransformConfig": + read_is_dlo = isinstance(self.permissions.read, DloPermission) + write_is_dlo = isinstance(self.permissions.write, DloPermission) + if self.is_streaming: + if not write_is_dlo: + raise ValueError( + "A streaming transform must write to a DLO " + "(permissions.write must be a 'dlo' entry)." + ) + elif read_is_dlo != write_is_dlo: + raise ValueError( + "permissions.read and permissions.write must both reference " + "DLOs or both reference DMOs (got " + f"read={type(self.permissions.read).__name__}, " + f"write={type(self.permissions.write).__name__})" + ) + return self class FunctionConfig(BaseConfig): @@ -402,23 +438,15 @@ class DmoPermission(BaseModel): dmo: list[str] +class StreamingSource(BaseModel): + type: str + name: str + + class Permissions(BaseModel): read: Union[DloPermission, DmoPermission] write: Union[DloPermission, DmoPermission] - @model_validator(mode="after") - def _no_mixed_layers(self) -> "Permissions": - read_is_dlo = isinstance(self.read, DloPermission) - write_is_dlo = isinstance(self.write, DloPermission) - if read_is_dlo != write_is_dlo: - raise ValueError( - "permissions.read and permissions.write must both reference " - "DLOs or both reference DMOs (got " - f"read={type(self.read).__name__}, " - f"write={type(self.write).__name__})" - ) - return self - def _permission_entries(perm: Union[DloPermission, DmoPermission]) -> list[str]: """Return the list of object names regardless of layer (DLO or DMO).""" @@ -490,7 +518,9 @@ def create_data_transform( ) -> dict: """Create a data transform in the DataCloud.""" script_name = metadata.name - request_hydrated = DATA_TRANSFORM_REQUEST_TEMPLATE.copy() + # Deep copy: the template's nested nodes/sources/macros dicts would + # otherwise be shared across calls and accumulate entries between deploys. + request_hydrated = copy.deepcopy(DATA_TRANSFORM_REQUEST_TEMPLATE) # Add nodes for each write entry (DLO or DMO) for i, name in enumerate( @@ -533,7 +563,7 @@ def create_data_transform( "definition": definition, "label": f"{metadata.name}", "name": f"{metadata.name}", - "type": "BATCH", + "type": "STREAMING" if data_transform_config.is_streaming else "BATCH", "dataSpaceName": data_transform_config.dataspace, } @@ -616,9 +646,18 @@ def deploy_full( callback=None, ) -> AccessTokenResponse: """Deploy a data transform in the DataCloud.""" + from datacustomcode.constants import SCRIPT_USE_IN_FEATURE_STREAMING + # prepare payload config = get_config(directory) + if ( + isinstance(config, DataTransformConfig) + and config.is_streaming + and not metadata.invokeOptions + ): + metadata.invokeOptions = [SCRIPT_USE_IN_FEATURE_STREAMING] + # create deployment and upload payload deployment = create_deployment(access_token, metadata) zip(directory, docker_network, metadata.codeType) diff --git a/src/datacustomcode/scan.py b/src/datacustomcode/scan.py index 5e50c5d..cdb6434 100644 --- a/src/datacustomcode/scan.py +++ b/src/datacustomcode/scan.py @@ -15,6 +15,7 @@ from __future__ import annotations import ast +import copy import json import os import sys @@ -22,6 +23,7 @@ Any, ClassVar, Dict, + Optional, Set, Union, ) @@ -32,6 +34,8 @@ from datacustomcode.version import get_version DATA_ACCESS_METHODS = ["read_dlo", "read_dmo", "write_to_dlo", "write_to_dmo"] +STREAMING_READ_METHODS = ["read_dlo_deltas", "read_dmo_deltas"] +STREAMING_WRITE_METHODS = ["write_dlo_deltas"] DATA_TRANSFORM_CONFIG_TEMPLATE = { "sdkVersion": get_version(), @@ -43,6 +47,20 @@ }, } +STREAMING_TRANSFORM_CONFIG_TEMPLATE = { + "sdkVersion": get_version(), + "entryPoint": "", + "dataspace": "default", + "streamingSource": { + "type": "dlo", + "name": "", + }, + "permissions": { + "read": {}, + "write": {}, + }, +} + FUNCTION_CONFIG_TEMPLATE = { "entryPoint": "", } @@ -160,6 +178,35 @@ def output_str(self) -> str: return next(iter(self.write_to_dmo)) +class StreamingDataAccessLayerCalls(pydantic.BaseModel): + read_dlo_deltas: bool + read_dmo_deltas: bool + write_dlo_deltas: frozenset[str] + + @pydantic.model_validator(mode="after") + def validate_access_layer(self) -> StreamingDataAccessLayerCalls: + if self.read_dlo_deltas and self.read_dmo_deltas: + raise ValueError( + "Cannot read DLO and DMO deltas in the same streaming transform." + ) + if not self.read_dlo_deltas and not self.read_dmo_deltas: + raise ValueError( + "A streaming transform must read from at least one DLO or DMO " + "delta stream (read_dlo_deltas / read_dmo_deltas)." + ) + if not self.write_dlo_deltas: + raise ValueError( + "A streaming transform must write to at least one DLO via " + "write_dlo_deltas." + ) + return self + + @property + def read_layer(self) -> str: + """Return the read source layer, ``"dlo"`` or ``"dmo"``.""" + return "dlo" if self.read_dlo_deltas else "dmo" + + class ClientMethodVisitor(ast.NodeVisitor): """AST Visitor that finds all instances of Client read/write method calls.""" @@ -168,6 +215,9 @@ def __init__(self) -> None: self._read_dmo_instances: set[str] = set() self._write_to_dlo_instances: set[str] = set() self._write_to_dmo_instances: set[str] = set() + self._read_dlo_deltas: bool = False + self._read_dmo_deltas: bool = False + self._write_dlo_deltas_instances: set[str] = set() self.variable_values: Dict[str, Union[str, None]] = {} def visit_Assign(self, node: ast.Assign) -> None: @@ -189,14 +239,15 @@ def visit_Call(self, node: ast.Call) -> None: node.func.value, ast.Name ): method_name = node.func.attr + + if method_name == "read_dlo_deltas": + self._read_dlo_deltas = True + elif method_name == "read_dmo_deltas": + self._read_dmo_deltas = True + if method_name in DATA_ACCESS_METHODS and node.args: arg = node.args[0] - name = None - - if isinstance(arg, ast.Constant) and isinstance(arg.value, str): - name = arg.value - elif isinstance(arg, ast.Name) and arg.id in self.variable_values: - name = self.variable_values[arg.id] + name = self._resolve_name_arg(arg) if name: if method_name == "read_dlo": @@ -207,8 +258,29 @@ def visit_Call(self, node: ast.Call) -> None: self._write_to_dlo_instances.add(name) elif method_name == "write_to_dmo": self._write_to_dmo_instances.add(name) + elif method_name in STREAMING_WRITE_METHODS and node.args: + name = self._resolve_name_arg(node.args[0]) + if name and method_name == "write_dlo_deltas": + self._write_dlo_deltas_instances.add(name) self.generic_visit(node) + def _resolve_name_arg(self, arg: ast.expr) -> Union[str, None]: + """Resolve a string-literal or tracked-variable first argument.""" + if isinstance(arg, ast.Constant) and isinstance(arg.value, str): + return arg.value + if isinstance(arg, ast.Name) and arg.id in self.variable_values: + return self.variable_values[arg.id] + return None + + @property + def is_streaming(self) -> bool: + """Whether any streaming (delta) access method was found.""" + return ( + self._read_dlo_deltas + or self._read_dmo_deltas + or bool(self._write_dlo_deltas_instances) + ) + def found(self) -> DataAccessLayerCalls: return DataAccessLayerCalls( read_dlo=frozenset(self._read_dlo_instances), @@ -217,6 +289,13 @@ def found(self) -> DataAccessLayerCalls: write_to_dmo=frozenset(self._write_to_dmo_instances), ) + def found_streaming(self) -> StreamingDataAccessLayerCalls: + return StreamingDataAccessLayerCalls( + read_dlo_deltas=self._read_dlo_deltas, + read_dmo_deltas=self._read_dmo_deltas, + write_dlo_deltas=frozenset(self._write_dlo_deltas_instances), + ) + class ImportVisitor(ast.NodeVisitor): """AST Visitor that extracts external package imports from Python code.""" @@ -301,23 +380,51 @@ def write_requirements_file(file_path: str) -> str: return requirements_path -def scan_file(file_path: str) -> DataAccessLayerCalls: - """Scan a single Python file for Client read/write method calls.""" +def _visit_file(file_path: str) -> ClientMethodVisitor: + """Parse a Python file and return the populated method visitor.""" with open(file_path, "r") as f: - code = f.read() - tree = ast.parse(code) - visitor = ClientMethodVisitor() - visitor.visit(tree) - return visitor.found() + tree = ast.parse(f.read()) + visitor = ClientMethodVisitor() + visitor.visit(tree) + return visitor + + +def scan_file(file_path: str) -> DataAccessLayerCalls: + """Scan a single Python file for batch Client read/write method calls.""" + return _visit_file(file_path).found() + + +def scan_file_streaming(file_path: str) -> StreamingDataAccessLayerCalls: + """Scan a single Python file for StreamingClient delta method calls.""" + return _visit_file(file_path).found_streaming() + +def file_is_streaming(file_path: str) -> bool: + """Return whether the entrypoint uses streaming (delta) access methods.""" + return _visit_file(file_path).is_streaming -def dc_config_json_from_file(file_path: str, type: str) -> dict[str, Any]: - """Create a Data Cloud Custom Code config JSON from a script.""" + +def dc_config_json_from_file( + file_path: str, type: str, streaming: bool = False +) -> dict[str, Any]: + """Create a Data Cloud Custom Code config JSON from a script. + + Args: + file_path: Path to the entrypoint. + type: Package type, ``"script"`` or ``"function"``. + streaming: For scripts, a streaming + (``streamingSource``) config instead of a batch one. + """ config: dict[str, Any] if type == "script": - config = DATA_TRANSFORM_CONFIG_TEMPLATE.copy() + template = ( + STREAMING_TRANSFORM_CONFIG_TEMPLATE + if streaming + else DATA_TRANSFORM_CONFIG_TEMPLATE + ) + config = copy.deepcopy(template) elif type == "function": - config = FUNCTION_CONFIG_TEMPLATE.copy() + config = copy.deepcopy(FUNCTION_CONFIG_TEMPLATE) config["entryPoint"] = os.path.basename(file_path) return config @@ -372,22 +479,53 @@ def update_config(file_path: str) -> dict[str, Any]: if package_type == "script": existing_config["dataspace"] = get_dataspace(existing_config) - output = scan_file(file_path) - read: dict[str, list[str]] = {} - if output.read_dlo: - read["dlo"] = list(output.read_dlo) - else: - read["dmo"] = list(output.read_dmo) - write: dict[str, list[str]] = {} - if output.write_to_dlo: - write["dlo"] = list(output.write_to_dlo) + if file_is_streaming(file_path): + _update_streaming_config(existing_config, file_path) else: - write["dmo"] = list(output.write_to_dmo) - - existing_config["permissions"] = {"read": read, "write": write} + existing_config.pop("streamingSource", None) + output = scan_file(file_path) + read: dict[str, list[str]] = {} + if output.read_dlo: + read["dlo"] = list(output.read_dlo) + else: + read["dmo"] = list(output.read_dmo) + write: dict[str, list[str]] = {} + if output.write_to_dlo: + write["dlo"] = list(output.write_to_dlo) + else: + write["dmo"] = list(output.write_to_dmo) + + existing_config["permissions"] = {"read": read, "write": write} return existing_config +def _update_streaming_config( + existing_config: dict[str, Any], file_path: str +) -> None: + output = scan_file_streaming(file_path) + read_layer = output.read_layer + + source = existing_config.get("streamingSource") + if not isinstance(source, dict): + source = {} + source_name = source.get("name", "") + existing_config["streamingSource"] = {"type": read_layer, "name": source_name} + + if not source_name: + logger.warning( + "streamingSource.name is empty in config.json. A streaming " + "transform must declare its read source; set streamingSource.name " + "to the DLO/DMO the transform reads from." + ) + + read_names = [source_name] if source_name else [] + write_names = list(output.write_dlo_deltas) + existing_config["permissions"] = { + "read": {read_layer: read_names}, + "write": {"dlo": write_names}, + } + + def get_dataspace(existing_config: dict[str, str]) -> str: if "dataspace" in existing_config: dataspace_value = existing_config["dataspace"] diff --git a/src/datacustomcode/template.py b/src/datacustomcode/template.py index 6807510..a543575 100644 --- a/src/datacustomcode/template.py +++ b/src/datacustomcode/template.py @@ -23,8 +23,12 @@ script_template_dir = os.path.join(os.path.dirname(__file__), "templates", "script") function_template_dir = os.path.join(os.path.dirname(__file__), "templates", "function") +STREAMING_EXAMPLE_ENTRYPOINT = os.path.join( + script_template_dir, "examples", "streaming_deltas", "entrypoint.py" +) -def copy_script_template(target_dir: str) -> None: + +def copy_script_template(target_dir: str, streaming: bool = False) -> None: """Copy the template to the target directory.""" os.makedirs(target_dir, exist_ok=True) @@ -39,6 +43,14 @@ def copy_script_template(target_dir: str) -> None: logger.debug(f"Copying file {source} to {destination}...") shutil.copy2(source, destination) + if streaming: + destination = os.path.join(target_dir, "payload", "entrypoint.py") + logger.debug( + f"Copying streaming example {STREAMING_EXAMPLE_ENTRYPOINT} to " + f"{destination}..." + ) + shutil.copy2(STREAMING_EXAMPLE_ENTRYPOINT, destination) + def copy_function_template(target_dir: str, use_in_feature: Optional[str]) -> None: os.makedirs(target_dir, exist_ok=True) diff --git a/tests/test_cli.py b/tests/test_cli.py index 7765560..ad55810 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -46,11 +46,14 @@ def test_init_command( result = runner.invoke(init, ["test_dir", "--code-type", "script"]) assert result.exit_code == 0 - mock_copy.assert_called_once_with("test_dir") + # A script with no --use-in-feature defaults to batch (streaming=False). + mock_copy.assert_called_once_with("test_dir", streaming=False) # Verify SDK config was written mock_write_sdk.assert_called_once_with("test_dir", {"type": "script"}) mock_scan.assert_called_once_with( - os.path.join("test_dir", "payload", "entrypoint.py"), "script" + os.path.join("test_dir", "payload", "entrypoint.py"), + "script", + streaming=False, ) mock_update.assert_called_once_with( os.path.join("test_dir", "payload", "entrypoint.py") @@ -69,6 +72,96 @@ def test_init_command( expected_content = json.dumps(mock_update.return_value, indent=2) assert expected_content in written_content + @patch("datacustomcode.template.copy_script_template") + @patch("datacustomcode.scan.update_config") + @patch("datacustomcode.scan.dc_config_json_from_file") + @patch("datacustomcode.scan.write_sdk_config") + @patch("builtins.open", new_callable=mock_open) + def test_init_command_streaming( + self, mock_file, mock_write_sdk, mock_scan, mock_update, mock_copy + ): + """Test init command with --use-in-feature StreamingTransform.""" + mock_scan.return_value = {"streamingSource": {"type": "dlo", "name": ""}} + mock_update.return_value = {"streamingSource": {"type": "dlo", "name": ""}} + + runner = CliRunner() + with runner.isolated_filesystem(): + os.makedirs(os.path.join("test_dir", "payload"), exist_ok=True) + + result = runner.invoke( + init, + [ + "test_dir", + "--code-type", + "script", + "--use-in-feature", + "StreamingTransform", + ], + ) + + assert result.exit_code == 0 + mock_copy.assert_called_once_with("test_dir", streaming=True) + mock_scan.assert_called_once_with( + os.path.join("test_dir", "payload", "entrypoint.py"), + "script", + streaming=True, + ) + + @patch("datacustomcode.template.copy_script_template") + @patch("datacustomcode.scan.update_config") + @patch("datacustomcode.scan.dc_config_json_from_file") + @patch("datacustomcode.scan.write_sdk_config") + @patch("builtins.open", new_callable=mock_open) + def test_init_command_batch_explicit( + self, mock_file, mock_write_sdk, mock_scan, mock_update, mock_copy + ): + """Test init command with explicit --use-in-feature BatchTransform.""" + mock_scan.return_value = {"permissions": {"read": {}, "write": {}}} + mock_update.return_value = {"permissions": {"read": {}, "write": {}}} + + runner = CliRunner() + with runner.isolated_filesystem(): + os.makedirs(os.path.join("test_dir", "payload"), exist_ok=True) + + result = runner.invoke( + init, + [ + "test_dir", + "--code-type", + "script", + "-u", + "BatchTransform", + ], + ) + + assert result.exit_code == 0 + mock_copy.assert_called_once_with("test_dir", streaming=False) + mock_scan.assert_called_once_with( + os.path.join("test_dir", "payload", "entrypoint.py"), + "script", + streaming=False, + ) + + def test_init_command_invalid_use_in_feature(self): + """A bad --use-in-feature value for a script aborts with an error.""" + runner = CliRunner() + with runner.isolated_filesystem(): + os.makedirs(os.path.join("test_dir", "payload"), exist_ok=True) + + result = runner.invoke( + init, + [ + "test_dir", + "--code-type", + "script", + "--use-in-feature", + "NotARealOption", + ], + ) + + assert result.exit_code != 0 + assert "Invalid --use-in-feature" in result.output + class TestDeploy: @patch("datacustomcode.deploy.deploy_full") diff --git a/tests/test_deploy.py b/tests/test_deploy.py index af804d3..9c203b8 100644 --- a/tests/test_deploy.py +++ b/tests/test_deploy.py @@ -17,6 +17,7 @@ DloPermission, DmoPermission, Permissions, + StreamingSource, get_config, ) @@ -671,6 +672,56 @@ def test_create_deployment_function_invoke_options(self, mock_make_api_call): assert isinstance(result, CreateDeploymentResponse) assert result.fileUploadUrl == "https://upload.example.com" + @patch("datacustomcode.deploy._make_api_call") + def test_create_deployment_default_path_no_invoke_options( + self, mock_make_api_call + ): + """Without invokeOptions, the deployment uses the default v63.0 path.""" + access_token = AccessTokenResponse( + access_token="test_token", instance_url="https://instance.example.com" + ) + metadata = CodeExtensionMetadata( + name="test_job", + version="1.0.0", + description="Test job", + computeType="CPU_M", + codeType="script", + ) + mock_make_api_call.return_value = { + "fileUploadUrl": "https://upload.example.com" + } + + create_deployment(access_token, metadata) + + url = mock_make_api_call.call_args[0][0] + assert "v63.0" in url + body = mock_make_api_call.call_args[1]["json"] + assert "invokeOptions" not in body + + @patch("datacustomcode.deploy._make_api_call") + def test_create_deployment_invoke_options_routes_to_v67(self, mock_make_api_call): + access_token = AccessTokenResponse( + access_token="test_token", instance_url="https://instance.example.com" + ) + metadata = CodeExtensionMetadata( + name="test_job", + version="1.0.0", + description="Test job", + computeType="CPU_M", + codeType="script", + invokeOptions=["StreamingTransform"], + ) + mock_make_api_call.return_value = { + "fileUploadUrl": "https://upload.example.com" + } + + create_deployment(access_token, metadata) + + url = mock_make_api_call.call_args[0][0] + assert "v67.0" in url + body = mock_make_api_call.call_args[1]["json"] + assert body["invokeOptions"] == ["StreamingTransform"] + class TestZip: @patch("datacustomcode.deploy.has_nonempty_requirements_file") @@ -1353,6 +1404,151 @@ def test_create_data_transform_dmo_missing_data_objects_raises( "/test/dir", access_token, metadata, data_transform_config ) + @patch("datacustomcode.deploy.get_config") + @patch("datacustomcode.deploy._make_api_call") + def test_create_data_transform_batch_type( + self, mock_make_api_call, mock_get_config + ): + """A non-streaming transform sends type BATCH.""" + access_token = AccessTokenResponse( + access_token="test_token", instance_url="https://instance.example.com" + ) + metadata = CodeExtensionMetadata( + name="batch_job", + version="1.0.0", + description="Batch job", + computeType="CPU_M", + codeType="script", + ) + data_transform_config = DataTransformConfig( + sdkVersion="1.0.0", + entryPoint="entrypoint.py", + dataspace="test_dataspace", + permissions=Permissions( + read=DloPermission(dlo=["input_dlo"]), + write=DloPermission(dlo=["output_dlo"]), + ), + ) + mock_make_api_call.return_value = {"id": "transform_id"} + + create_data_transform( + "/test/dir", access_token, metadata, data_transform_config + ) + + request_body = mock_make_api_call.call_args[1]["json"] + assert request_body["type"] == "BATCH" + + @patch("datacustomcode.deploy.get_config") + @patch("datacustomcode.deploy._make_api_call") + def test_create_data_transform_streaming_type( + self, mock_make_api_call, mock_get_config + ): + """A streaming (streamingSource) transform sends type STREAMING.""" + access_token = AccessTokenResponse( + access_token="test_token", instance_url="https://instance.example.com" + ) + metadata = CodeExtensionMetadata( + name="streaming_job", + version="1.0.0", + description="Streaming job", + computeType="CPU_M", + codeType="script", + ) + data_transform_config = DataTransformConfig( + sdkVersion="1.0.0", + entryPoint="entrypoint.py", + dataspace="test_dataspace", + streamingSource=StreamingSource(type="dlo", name="input_dlo"), + permissions=Permissions( + read=DloPermission(dlo=["input_dlo"]), + write=DloPermission(dlo=["output_dlo"]), + ), + ) + mock_make_api_call.return_value = {"id": "transform_id"} + + create_data_transform( + "/test/dir", access_token, metadata, data_transform_config + ) + + request_body = mock_make_api_call.call_args[1]["json"] + assert request_body["type"] == "STREAMING" + manifest = request_body["definition"]["manifest"] + assert manifest["sources"] == {"source1": {"relation_name": "input_dlo"}} + assert manifest["nodes"]["node1"]["relation_name"] == "output_dlo" + + +class TestDataTransformConfigStreaming: + """The streamingSource field and its effect on layer validation.""" + + def test_is_streaming_true_when_source_present(self): + config = DataTransformConfig( + sdkVersion="1.0.0", + entryPoint="entrypoint.py", + dataspace="default", + streamingSource=StreamingSource(type="dlo", name="input_dlo"), + permissions=Permissions( + read=DloPermission(dlo=["input_dlo"]), + write=DloPermission(dlo=["output_dlo"]), + ), + ) + assert config.is_streaming is True + + def test_is_streaming_false_when_source_absent(self): + config = DataTransformConfig( + sdkVersion="1.0.0", + entryPoint="entrypoint.py", + dataspace="default", + permissions=Permissions( + read=DloPermission(dlo=["input_dlo"]), + write=DloPermission(dlo=["output_dlo"]), + ), + ) + assert config.is_streaming is False + + def test_streaming_allows_dmo_read_dlo_write(self): + """Streaming may read a DMO change feed and write a DLO (writes are + DLO-only)""" + config = DataTransformConfig( + sdkVersion="1.0.0", + entryPoint="entrypoint.py", + dataspace="default", + streamingSource=StreamingSource(type="dmo", name="input_dmo__dlm"), + permissions=Permissions( + read=DmoPermission(dmo=["input_dmo__dlm"]), + write=DloPermission(dlo=["output_dlo"]), + ), + ) + assert config.is_streaming is True + assert isinstance(config.permissions.read, DmoPermission) + assert isinstance(config.permissions.write, DloPermission) + + def test_streaming_rejects_dmo_write(self): + """Streaming writes must target a DLO.""" + with pytest.raises(ValueError, match="must write to a DLO"): + DataTransformConfig( + sdkVersion="1.0.0", + entryPoint="entrypoint.py", + dataspace="default", + streamingSource=StreamingSource(type="dlo", name="input_dlo"), + permissions=Permissions( + read=DloPermission(dlo=["input_dlo"]), + write=DmoPermission(dmo=["output_dmo__dlm"]), + ), + ) + + def test_batch_still_rejects_mixed_layers(self): + """Without a streamingSource, mixed read/write layers are still rejected.""" + with pytest.raises(ValueError, match="both reference"): + DataTransformConfig( + sdkVersion="1.0.0", + entryPoint="entrypoint.py", + dataspace="default", + permissions=Permissions( + read=DloPermission(dlo=["input_dlo"]), + write=DmoPermission(dmo=["output_dmo__dlm"]), + ), + ) + class TestDeployFull: @patch("datacustomcode.deploy.get_config") @@ -1466,6 +1662,96 @@ def test_deploy_full_client_credentials( ) assert result == access_token + @patch("datacustomcode.deploy.get_config") + @patch("datacustomcode.deploy.create_data_transform") + @patch("datacustomcode.deploy.wait_for_deployment") + @patch("datacustomcode.deploy.upload_zip") + @patch("datacustomcode.deploy.zip") + @patch("datacustomcode.deploy.create_deployment") + def test_deploy_full_streaming_sets_invoke_options( + self, + mock_create_deployment, + mock_zip, + mock_upload_zip, + mock_wait, + mock_create_transform, + mock_get_config, + ): + """A streaming config makes deploy_full set invokeOptions on the + metadata before creating the deployment.""" + data_transform_config = DataTransformConfig( + sdkVersion="1.0.0", + entryPoint="entrypoint.py", + dataspace="test_dataspace", + streamingSource=StreamingSource(type="dlo", name="input_dlo"), + permissions=Permissions( + read=DloPermission(dlo=["input_dlo"]), + write=DloPermission(dlo=["output_dlo"]), + ), + ) + mock_get_config.return_value = data_transform_config + metadata = CodeExtensionMetadata( + name="test_job", + version="1.0.0", + description="Test job", + computeType="CPU_M", + codeType="script", + ) + access_token = AccessTokenResponse( + access_token="test_token", instance_url="https://instance.example.com" + ) + mock_create_deployment.return_value = CreateDeploymentResponse( + fileUploadUrl="https://upload.example.com" + ) + + deploy_full("/test/dir", metadata, access_token, "default") + + assert metadata.invokeOptions == ["StreamingTransform"] + + @patch("datacustomcode.deploy.get_config") + @patch("datacustomcode.deploy.create_data_transform") + @patch("datacustomcode.deploy.wait_for_deployment") + @patch("datacustomcode.deploy.upload_zip") + @patch("datacustomcode.deploy.zip") + @patch("datacustomcode.deploy.create_deployment") + def test_deploy_full_batch_leaves_invoke_options_unset( + self, + mock_create_deployment, + mock_zip, + mock_upload_zip, + mock_wait, + mock_create_transform, + mock_get_config, + ): + """A batch config must not set invokeOptions (stays on the v63.0 path).""" + data_transform_config = DataTransformConfig( + sdkVersion="1.0.0", + entryPoint="entrypoint.py", + dataspace="test_dataspace", + permissions=Permissions( + read=DloPermission(dlo=["input_dlo"]), + write=DloPermission(dlo=["output_dlo"]), + ), + ) + mock_get_config.return_value = data_transform_config + metadata = CodeExtensionMetadata( + name="test_job", + version="1.0.0", + description="Test job", + computeType="CPU_M", + codeType="script", + ) + access_token = AccessTokenResponse( + access_token="test_token", instance_url="https://instance.example.com" + ) + mock_create_deployment.return_value = CreateDeploymentResponse( + fileUploadUrl="https://upload.example.com" + ) + + deploy_full("/test/dir", metadata, access_token, "default") + + assert metadata.invokeOptions is None + class TestRunDataTransform: @patch("datacustomcode.deploy._make_api_call") diff --git a/tests/test_scan.py b/tests/test_scan.py index 2acbc25..716c233 100644 --- a/tests/test_scan.py +++ b/tests/test_scan.py @@ -12,8 +12,10 @@ SDK_CONFIG_FILE, DataAccessLayerCalls, dc_config_json_from_file, + file_is_streaming, scan_file, scan_file_for_imports, + scan_file_streaming, update_config, write_requirements_file, write_sdk_config, @@ -339,6 +341,27 @@ def test_dlo_to_dlo_config(self): os.remove(sdk_config_path) os.rmdir(os.path.dirname(sdk_config_path)) + def test_streaming_config_scaffolds_streaming_source(self): + """dc_config_json_from_file(streaming=True) scaffolds streamingSource.""" + temp_path = create_test_script(STREAMING_DLO_ENTRYPOINT) + try: + result = dc_config_json_from_file(temp_path, "script", streaming=True) + assert result["entryPoint"] == os.path.basename(temp_path) + assert result["dataspace"] == "default" + assert result["sdkVersion"] == get_version() + assert result["streamingSource"] == {"type": "dlo", "name": ""} + finally: + os.unlink(temp_path) + + def test_batch_config_has_no_streaming_source(self): + """The default (batch) script template has no streamingSource.""" + temp_path = create_test_script("x = 1\n") + try: + result = dc_config_json_from_file(temp_path, "script") + assert "streamingSource" not in result + finally: + os.unlink(temp_path) + def test_dmo_to_dmo_config(self): """Test generating config JSON for DMO to DMO operations.""" content = textwrap.dedent( @@ -694,6 +717,190 @@ def my_function(event, context): os.rmdir(os.path.dirname(sdk_config_path)) +STREAMING_DLO_ENTRYPOINT = textwrap.dedent( + """ + from datacustomcode.client import StreamingClient + + client = StreamingClient() + deltas = client.read_dlo_deltas() + transformed = deltas.withColumn("x", deltas.x) + query = client.write_dlo_deltas("Account_copy__dll", transformed) + query.awaitTermination() + """ +) + +STREAMING_DMO_ENTRYPOINT = textwrap.dedent( + """ + from datacustomcode.client import StreamingClient + + client = StreamingClient() + deltas = client.read_dmo_deltas() + query = client.write_dlo_deltas("Account_copy__dll", deltas) + query.awaitTermination() + """ +) + + +class TestStreamingScan: + """Tests for streaming (delta) detection and extraction.""" + + def test_file_is_streaming_true_for_delta_methods(self): + temp_path = create_test_script(STREAMING_DLO_ENTRYPOINT) + try: + assert file_is_streaming(temp_path) is True + finally: + os.unlink(temp_path) + + def test_file_is_streaming_false_for_batch(self): + content = textwrap.dedent( + """ + from datacustomcode.client import Client + + client = Client() + df = client.read_dlo("input_dlo") + client.write_to_dlo("output_dlo", df, "overwrite") + """ + ) + temp_path = create_test_script(content) + try: + assert file_is_streaming(temp_path) is False + finally: + os.unlink(temp_path) + + def test_scan_file_streaming_dlo(self): + temp_path = create_test_script(STREAMING_DLO_ENTRYPOINT) + try: + result = scan_file_streaming(temp_path) + assert result.read_dlo_deltas is True + assert result.read_dmo_deltas is False + assert result.read_layer == "dlo" + assert "Account_copy__dll" in result.write_dlo_deltas + finally: + os.unlink(temp_path) + + def test_scan_file_streaming_dmo_read(self): + temp_path = create_test_script(STREAMING_DMO_ENTRYPOINT) + try: + result = scan_file_streaming(temp_path) + assert result.read_dmo_deltas is True + assert result.read_layer == "dmo" + assert "Account_copy__dll" in result.write_dlo_deltas + finally: + os.unlink(temp_path) + + def test_scan_file_streaming_requires_read(self): + content = textwrap.dedent( + """ + from datacustomcode.client import StreamingClient + + client = StreamingClient() + client.write_dlo_deltas("Account_copy__dll", some_df) + """ + ) + temp_path = create_test_script(content) + try: + with pytest.raises(ValueError, match="at least one DLO or DMO delta"): + scan_file_streaming(temp_path) + finally: + os.unlink(temp_path) + + def test_scan_file_streaming_requires_write(self): + content = textwrap.dedent( + """ + from datacustomcode.client import StreamingClient + + client = StreamingClient() + deltas = client.read_dlo_deltas() + """ + ) + temp_path = create_test_script(content) + try: + with pytest.raises(ValueError, match="must write to at least one DLO"): + scan_file_streaming(temp_path) + finally: + os.unlink(temp_path) + + +class TestStreamingUpdateConfig: + """Tests for update_config on streaming entrypoints.""" + + def _run(self, entrypoint_src: str, initial_config: dict) -> dict: + temp_path = create_test_script(entrypoint_src) + file_dir = os.path.dirname(temp_path) + config_path = os.path.join(file_dir, "config.json") + sdk_config_path = create_sdk_config(file_dir, "script") + try: + with open(config_path, "w") as f: + json.dump(initial_config, f) + return update_config(temp_path) + finally: + os.remove(temp_path) + if os.path.exists(config_path): + os.remove(config_path) + if os.path.exists(sdk_config_path): + os.remove(sdk_config_path) + os.rmdir(os.path.dirname(sdk_config_path)) + + def test_preserves_streaming_source_name(self): + """config.json's streamingSource.name is authoritative and preserved.""" + updated = self._run( + STREAMING_DLO_ENTRYPOINT, + { + "sdkVersion": "1.0.0", + "entryPoint": "old.py", + "dataspace": "default", + "streamingSource": {"type": "dlo", "name": "Account_Home__dll"}, + "permissions": {"read": {}, "write": {}}, + }, + ) + assert updated["streamingSource"] == { + "type": "dlo", + "name": "Account_Home__dll", + } + assert updated["permissions"]["read"]["dlo"] == ["Account_Home__dll"] + assert updated["permissions"]["write"]["dlo"] == ["Account_copy__dll"] + + def test_reasserts_layer_from_code(self): + """A code read layer of DMO overrides a stale DLO type in config.""" + updated = self._run( + STREAMING_DMO_ENTRYPOINT, + { + "sdkVersion": "1.0.0", + "entryPoint": "old.py", + "dataspace": "default", + "streamingSource": {"type": "dlo", "name": "Account__dlm"}, + "permissions": {"read": {}, "write": {}}, + }, + ) + assert updated["streamingSource"]["type"] == "dmo" + assert updated["streamingSource"]["name"] == "Account__dlm" + assert updated["permissions"]["read"]["dmo"] == ["Account__dlm"] + + def test_batch_entrypoint_drops_stale_streaming_source(self): + """Switching a streaming entrypoint to batch clears streamingSource.""" + batch_src = textwrap.dedent( + """ + from datacustomcode.client import Client + + client = Client() + df = client.read_dlo("input_dlo") + client.write_to_dlo("output_dlo", df, "overwrite") + """ + ) + updated = self._run( + batch_src, + { + "sdkVersion": "1.0.0", + "entryPoint": "old.py", + "dataspace": "default", + "streamingSource": {"type": "dlo", "name": "stale__dll"}, + "permissions": {"read": {}, "write": {}}, + }, + ) + assert "streamingSource" not in updated + assert updated["permissions"]["read"]["dlo"] == ["input_dlo"] + + class TestDataAccessLayerCalls: """Tests for the DataAccessLayerCalls class directly.""" From 5738e807e4472c6ebc8d16fd7caef894d59f45a4 Mon Sep 17 00:00:00 2001 From: Joshua Catt Date: Thu, 23 Jul 2026 10:21:16 -0400 Subject: [PATCH 2/2] lint --- src/datacustomcode/scan.py | 5 +---- tests/test_deploy.py | 4 +--- 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/src/datacustomcode/scan.py b/src/datacustomcode/scan.py index cdb6434..059925a 100644 --- a/src/datacustomcode/scan.py +++ b/src/datacustomcode/scan.py @@ -23,7 +23,6 @@ Any, ClassVar, Dict, - Optional, Set, Union, ) @@ -499,9 +498,7 @@ def update_config(file_path: str) -> dict[str, Any]: return existing_config -def _update_streaming_config( - existing_config: dict[str, Any], file_path: str -) -> None: +def _update_streaming_config(existing_config: dict[str, Any], file_path: str) -> None: output = scan_file_streaming(file_path) read_layer = output.read_layer diff --git a/tests/test_deploy.py b/tests/test_deploy.py index 9c203b8..05063a4 100644 --- a/tests/test_deploy.py +++ b/tests/test_deploy.py @@ -673,9 +673,7 @@ def test_create_deployment_function_invoke_options(self, mock_make_api_call): assert result.fileUploadUrl == "https://upload.example.com" @patch("datacustomcode.deploy._make_api_call") - def test_create_deployment_default_path_no_invoke_options( - self, mock_make_api_call - ): + def test_create_deployment_default_path_no_invoke_options(self, mock_make_api_call): """Without invokeOptions, the deployment uses the default v63.0 path.""" access_token = AccessTokenResponse( access_token="test_token", instance_url="https://instance.example.com"