Several bug-fixes in deploy cmd, and added better starting point for entrypoint & config.json

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

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 from html import unescape
+import json
 import os
 import shutil
 import tarfile
@@ -30,11 +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
 
 if TYPE_CHECKING:
     from datacustomcode.credentials import Credentials
@@ -77,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)}"
@@ -223,77 +226,65 @@ 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
 
 
-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 Permissions(BaseModel):
+    read: Union[DloPermission]
+    write: Union[DloPermission]
 
 
-def verify_data_transform_config(directory: str) -> None:
-    """Verify that the data transform config.json file exists in the directory."""
+class DloPermission(BaseModel):
+    dlo: list[str]
+
+
+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")
-    if not os.path.exists(config_path):
-        raise FileNotFoundError(f"config.json not found in {directory}")
+    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 in {config_path}"
+        ) from err
+    except json.JSONDecodeError as err:
+        raise ValueError(
+            f"config.json in {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 in {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}")
 
 
@@ -306,28 +297,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)

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