Merge pull request #8 from forcedotcom/sdkVersion-config-2

Refactor config.json related parts of scan.py and deploy.py

Author: markdlv-sf

src/datacustomcode/cli.py

@@ -101,13 +101,25 @@ def deploy(profile: str, path: str, name: str, version: str, description: str):
 @cli.command()
 @click.argument("directory", default=".")
 def init(directory: str):
+    from datacustomcode.scan import dc_config_json_from_file
     from datacustomcode.template import copy_template
 
     click.echo("Copying template to " + click.style(directory, fg="blue", bold=True))
     copy_template(directory)
+    entrypoint_path = os.path.join(directory, "payload", "entrypoint.py")
+    config_location = os.path.join(os.path.dirname(entrypoint_path), "config.json")
+    config_json = dc_config_json_from_file(entrypoint_path)
+    with open(config_location, "w") as f:
+        json.dump(config_json, f, indent=2)
+
     click.echo(
         "Start developing by updating the code in "
-        + click.style(f"{directory}/payload/entrypoint.py", fg="blue", bold=True)
+        + click.style(entrypoint_path, fg="blue", bold=True)
+    )
+    click.echo(
+        "You can run "
+        + click.style(f"datacustomcode scan {entrypoint_path}", fg="blue", bold=True)
+        + " to automatically update config.json when you make changes to your code"
     )
 
 

src/datacustomcode/deploy.py

@@ -31,12 +31,11 @@
 )
 
 from loguru import logger
+import pydantic
 from pydantic import BaseModel
 import requests
 
 from datacustomcode.cmd import cmd_output
-from datacustomcode.scan import scan_file
-from datacustomcode.version import get_version
 
 if TYPE_CHECKING:
     from datacustomcode.credentials import Credentials
@@ -79,8 +78,10 @@ def _make_api_call(
     logger.debug(f"Request params: {kwargs}")
 
     response = requests.request(method=method, url=url, headers=headers, **kwargs)
-    response.raise_for_status()
     json_response = response.json()
+    if response.status_code >= 400:
+        logger.debug(f"Error Response: {json_response}")
+    response.raise_for_status()
     assert isinstance(
         json_response, dict
     ), f"Unexpected response type: {type(json_response)}"
@@ -225,89 +226,62 @@ def wait_for_deployment(
             callback(status)
         if status == "Deployed":
             logger.debug(
-                "Deployment completed, Elapsed time: {time.time() - start_time}"
+                f"Deployment completed.\nElapsed time: {time.time() - start_time}"
             )
             break
         time.sleep(1)
 
 
 DATA_TRANSFORM_REQUEST_TEMPLATE: dict[str, Any] = {
-    "metadata": {
-        "dbt_schema_version": "https://schemas.getdbt.com/dbt/manifest/v8.json",
-        "dbt_version": "1.4.6",
-        "generated_at": "2023-04-25T18:54:11.375589Z",
-        "invocation_id": "d6c68c69-533a-4d54-861e-1493d6cd8092",
-        "env": {},
-        "project_id": "jaffle_shop",
-        "user_id": "1ca8403c-a1a5-43af-8b88-9265e948b9d2",
-        "send_anonymous_usage_stats": True,
-        "adapter_type": "spark",
-    },
-    "nodes": {
-        "model.dcexample.dim_listings_w_hosts": {
-            "name": "dim_listings_w_hosts",
-            "resource_type": "model",
-            "relation_name": "{OUTPUT_DLO}",
-            "config": {"materialized": "table"},
-            "compiled_code": "",
-            "depends_on": {"nodes": []},
-        }
-    },
-    "sources": {
-        "source.dcexample.listings": {
-            "name": "listings",
-            "resource_type": "source",
-            "relation_name": "{INPUT_DLO}",
-            "identifier": "{INPUT_DLO}",
-        }
-    },
+    "nodes": {},
+    "sources": {},
     "macros": {
-        "macro.dcexample.byoc": {
-            "name": "byoc_example",
-            "resource_type": "macro",
-            "path": "",
-            "original_file_path": "",
-            "unique_id": "unique id",
-            "macro_sql": "",
-            "supported_languages": None,
+        "macro.byoc": {
             "arguments": [{"name": "{SCRIPT_NAME}", "type": "BYOC_SCRIPT"}],
         }
     },
 }
 
 
 class DataTransformConfig(BaseModel):
-    input: Union[str, list[str]]
-    output: Union[str, list[str]]
+    sdkVersion: str
+    entryPoint: str
+    dataspace: str
+    permissions: Permissions
 
 
-DATA_TRANSFORM_CONFIG_TEMPLATE: dict[str, Any] = {
-    "entryPoint": "entrypoint.py",
-    "dataspace": "default",
-    "permissions": {"read": {"dlo": ""}, "write": {"dlo": ""}},
-    "sdkVersion": get_version(),
-}
+class Permissions(BaseModel):
+    read: Union[DloPermission]
+    write: Union[DloPermission]
 
 
-def get_data_transform_config(directory: str) -> DataTransformConfig:
-    """Get the data transform config from the entrypoint.py file."""
-    entrypoint_file = os.path.join(directory, "entrypoint.py")
-    data_access_layer_calls = scan_file(entrypoint_file)
-    input_ = data_access_layer_calls.input_str
-    output = data_access_layer_calls.output_str
-    return DataTransformConfig(input=input_, output=output)
+class DloPermission(BaseModel):
+    dlo: list[str]
 
 
-def create_data_transform_config(directory: str) -> None:
-    """Create a data transform config.json file in the directory."""
-    data_transform_config = get_data_transform_config(directory)
-    request_hydrated = DATA_TRANSFORM_CONFIG_TEMPLATE.copy()
-    request_hydrated["permissions"]["read"]["dlo"] = data_transform_config.input
-    request_hydrated["permissions"]["write"]["dlo"] = data_transform_config.output
-    logger.debug(f"Creating data transform config in {directory}")
-    json.dump(
-        request_hydrated, open(os.path.join(directory, "config.json"), "w"), indent=4
-    )
+def get_data_transform_config(directory: str) -> DataTransformConfig:
+    """Get the data transform config from the config.json file."""
+    config_path = os.path.join(directory, "config.json")
+    try:
+        with open(config_path, "r") as f:
+            config = json.loads(f.read())
+            return DataTransformConfig(**config)
+    except FileNotFoundError as err:
+        raise FileNotFoundError(f"config.json not found at {config_path}") from err
+    except json.JSONDecodeError as err:
+        raise ValueError(f"config.json at {config_path} is not valid JSON") from err
+    except pydantic.ValidationError as err:
+        missing_fields = [str(err["loc"][0]) for err in err.errors()]
+        raise ValueError(
+            f"config.json at {config_path} is missing required "
+            f"fields: {', '.join(missing_fields)}"
+        ) from err
+
+
+def verify_data_transform_config(directory: str) -> None:
+    """Verify the data transform config.json contents."""
+    get_data_transform_config(directory)
+    logger.debug(f"Verified data transform config in {directory}")
 
 
 def create_data_transform(
@@ -319,28 +293,31 @@ def create_data_transform(
     script_name = metadata.name
     data_transform_config = get_data_transform_config(directory)
     request_hydrated = DATA_TRANSFORM_REQUEST_TEMPLATE.copy()
-    request_hydrated["nodes"]["model.dcexample.dim_listings_w_hosts"][
-        "relation_name"
-    ] = data_transform_config.input
-    request_hydrated["sources"]["source.dcexample.listings"][
-        "relation_name"
-    ] = data_transform_config.output
-    request_hydrated["sources"]["source.dcexample.listings"][
-        "identifier"
-    ] = data_transform_config.output
-    request_hydrated["macros"]["macro.dcexample.byoc"]["arguments"][0][
-        "name"
-    ] = script_name
+
+    # Add nodes for each write DLO
+    for i, dlo in enumerate(data_transform_config.permissions.write.dlo, 1):
+        request_hydrated["nodes"][f"node{i}"] = {
+            "relation_name": dlo,
+            "config": {"materialized": "table"},
+            "compiled_code": "",
+        }
+
+    # Add sources for each read DLO
+    for i, dlo in enumerate(data_transform_config.permissions.read.dlo, 1):
+        request_hydrated["sources"][f"source{i}"] = {"relation_name": dlo}
+
+    request_hydrated["macros"]["macro.byoc"]["arguments"][0]["name"] = script_name
 
     body = {
         "definition": {
-            "type": "DBT",
+            "type": "DCSQL",
             "manifest": request_hydrated,
             "version": "56.0",
         },
         "label": f"{metadata.name}",
         "name": f"{metadata.name}",
         "type": "BATCH",
+        "dataSpaceName": data_transform_config.dataspace,
     }
 
     url = _join_strip_url(access_token.instance_url, DATA_TRANSFORMS_PATH)
@@ -359,7 +336,7 @@ def deploy_full(
 
     # prepare payload
     prepare_dependency_archive(directory)
-    create_data_transform_config(directory)
+    verify_data_transform_config(directory)
 
     # create deployment and upload payload
     deployment = create_deployment(access_token, metadata)

src/datacustomcode/scan.py

@@ -16,15 +16,27 @@
 
 import ast
 from typing import (
+    Any,
     Dict,
-    List,
     Union,
 )
 
 import pydantic
 
+from datacustomcode.version import get_version
+
 DATA_ACCESS_METHODS = ["read_dlo", "read_dmo", "write_to_dlo", "write_to_dmo"]
 
+DATA_TRANSFORM_CONFIG_TEMPLATE = {
+    "sdkVersion": get_version(),
+    "entryPoint": "",
+    "dataspace": "default",
+    "permissions": {
+        "read": {},
+        "write": {},
+    },
+}
+
 
 class DataAccessLayerCalls(pydantic.BaseModel):
     read_dlo: frozenset[str]
@@ -129,25 +141,22 @@ def scan_file(file_path: str) -> DataAccessLayerCalls:
         return visitor.found()
 
 
-def dc_config_json_from_file(file_path: str) -> dict:
+def dc_config_json_from_file(file_path: str) -> dict[str, Any]:
     """Create a Data Cloud Custom Code config JSON from a script."""
     output = scan_file(file_path)
-    read = {}
+    config = DATA_TRANSFORM_CONFIG_TEMPLATE.copy()
+    config["entryPoint"] = file_path.rpartition("/")[-1]
+
+    read: dict[str, list[str]] = {}
     if output.read_dlo:
         read["dlo"] = list(output.read_dlo)
     else:
         read["dmo"] = list(output.read_dmo)
-    write = {}
+    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)
-    config: Dict[str, Union[str, Dict[str, Dict[str, List[str]]]]] = {
-        "entryPoint": file_path.rpartition("/")[-1],
-        "dataspace": "default",
-        "permissions": {
-            "read": read,
-            "write": write,
-        },
-    }
+
+    config["permissions"] = {"read": read, "write": write}
     return config

src/datacustomcode/templates/payload/entrypoint.py

@@ -1,9 +1,24 @@
+from pyspark.sql.functions import col, upper
+
 from datacustomcode.client import Client
+from datacustomcode.io.writer.base import WriteMode
 
 
 def main():
-    client = Client()  # noqa: F841
-    # TODO: Add your custom code here
+    client = Client()
+
+    df = client.read_dlo("Account_Home__dll")
+
+    # Perform transformations on the DataFrame
+    df_upper1 = df.withColumn("Description__c", upper(col("Description__c")))
+
+    # Drop specific columns related to relationships
+    df_upper1 = df_upper1.drop("KQ_ParentId__c")
+    df_upper1 = df_upper1.drop("KQ_Id__c")
+
+    # Save the transformed DataFrame
+    dlo_name = "Account_Home_copy__dll"
+    client.write_to_dlo(dlo_name, df_upper1, write_mode=WriteMode.APPEND)
 
 
 if __name__ == "__main__":

tests/test_cli.py

@@ -0,0 +1,49 @@
+import json
+import os
+from unittest.mock import mock_open, patch
+
+from click.testing import CliRunner
+
+from datacustomcode.cli import init
+
+
+class TestInit:
+    @patch("datacustomcode.template.copy_template")
+    @patch("datacustomcode.scan.dc_config_json_from_file")
+    @patch("builtins.open", new_callable=mock_open)
+    def test_init_command(self, mock_file, mock_scan, mock_copy):
+        """Test init command."""
+        mock_scan.return_value = {
+            "sdkVersion": "1.0.0",
+            "entryPoint": "entrypoint.py",
+            "dataspace": "default",
+            "permissions": {
+                "read": {"dlo": ["input_dlo"]},
+                "write": {"dlo": ["output_dlo"]},
+            },
+        }
+
+        runner = CliRunner()
+        with runner.isolated_filesystem():
+            # Create test directory structure
+            os.makedirs(os.path.join("test_dir", "payload"), exist_ok=True)
+
+            result = runner.invoke(init, ["test_dir"])
+
+            assert result.exit_code == 0
+            mock_copy.assert_called_once_with("test_dir")
+            mock_scan.assert_called_once_with(
+                os.path.join("test_dir", "payload", "entrypoint.py")
+            )
+
+            # Verify the config.json was written with the correct content
+            mock_file.assert_any_call(
+                os.path.join("test_dir", "payload", "config.json"), "w"
+            )
+
+            # Get all write calls and join them to get the complete written content
+            written_content = "".join(
+                call.args[0] for call in mock_file().write.call_args_list
+            )
+            expected_content = json.dumps(mock_scan.return_value, indent=2)
+            assert written_content == expected_content