salesforce-data-customcode 0.1.0__py3-none-any.whl

datacustomcode/io/writer/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

datacustomcode/io/writer/base.py

@@ -0,0 +1,49 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from abc import abstractmethod
+from enum import Enum
+from typing import TYPE_CHECKING
+
+from datacustomcode.io.base import BaseDataAccessLayer
+
+if TYPE_CHECKING:
+    from pyspark.sql import DataFrame as PySparkDataFrame, SparkSession
+
+
+class WriteMode(str, Enum):
+    APPEND = "append"
+    OVERWRITE = "overwrite"
+    OVERWRITE_PARTITIONS = "overwrite_partitions"
+    MERGE = "merge"
+    MERGE_UPSERT_DELETE = "merge_upsert_delete"
+
+
+class BaseDataCloudWriter(BaseDataAccessLayer):
+    """Base class for Data Cloud writers."""
+
+    def __init__(self, spark: SparkSession) -> None:
+        self.spark = spark
+
+    @abstractmethod
+    def write_to_dlo(
+        self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode
+    ) -> None: ...
+
+    @abstractmethod
+    def write_to_dmo(
+        self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode
+    ) -> None: ...

datacustomcode/io/writer/csv.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from pyspark.sql import DataFrame as PySparkDataFrame
+
+from datacustomcode.io.writer.base import BaseDataCloudWriter, WriteMode
+
+SUFFIX = ".csv"
+
+
+class CSVDataCloudWriter(BaseDataCloudWriter):
+    CONFIG_NAME = "CSVDataCloudWriter"
+
+    def write_to_dlo(
+        self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode
+    ) -> None:
+        # Only add the suffix if it's not already there
+        if not name.lower().endswith(SUFFIX):
+            name = f"{name}{SUFFIX}"
+        dataframe.write.csv(name, mode=write_mode)
+
+    def write_to_dmo(
+        self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode
+    ) -> None:
+        # Only add the suffix if it's not already there
+        if not name.lower().endswith(SUFFIX):
+            name = f"{name}{SUFFIX}"
+        dataframe.write.csv(name, mode=write_mode)

datacustomcode/io/writer/print.py

@@ -0,0 +1,33 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from pyspark.sql import DataFrame as PySparkDataFrame
+
+from datacustomcode.io.writer.base import BaseDataCloudWriter, WriteMode
+
+
+class PrintDataCloudWriter(BaseDataCloudWriter):
+    CONFIG_NAME = "PrintDataCloudWriter"
+
+    def write_to_dlo(
+        self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode
+    ) -> None:
+        dataframe.show()
+
+    def write_to_dmo(
+        self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode
+    ) -> None:
+        dataframe.show()

datacustomcode/mixin.py

@@ -0,0 +1,94 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from typing import ClassVar, TypeVar
+
+
+def _get_all_subclass_descendants(cls: type) -> list[type]:
+    all_subclasses = [cls]
+
+    for subclass in cls.__subclasses__():
+        all_subclasses.append(subclass)
+        all_subclasses.extend(_get_all_subclass_descendants(subclass))
+
+    return all_subclasses
+
+
+_V = TypeVar("_V", bound="UserExtendableNamedConfigMixin")
+
+
+class UserExtendableNamedConfigMixin:
+    """Allows users to access extended classes by name in config.
+
+    Our client code is text config driven. This means that if a user extends readers
+    and writers, they wouldn't be able to add their own classes if they want to use a
+    config. This allows them to name their subclasses and let them be found at runtime
+    by the config driven execution.
+    """
+
+    CONFIG_NAME: str
+
+    _registered_config_names: ClassVar[dict[str, type]] = {}
+
+    def __init_subclass__(cls, **kwargs):
+        """Validate CONFIG_NAME uniqueness across all subclasses."""
+        super().__init_subclass__(**kwargs)
+        if "CONFIG_NAME" not in cls.__dict__:
+            return
+        if cls.CONFIG_NAME is None or cls.CONFIG_NAME == "":
+            return
+
+        if cls.CONFIG_NAME in UserExtendableNamedConfigMixin._registered_config_names:
+            existing_class = UserExtendableNamedConfigMixin._registered_config_names[
+                cls.CONFIG_NAME
+            ]
+            raise TypeError(
+                f"Class {cls.__name__} has the same CONFIG_NAME ('{cls.CONFIG_NAME}') "
+                f"as existing class {existing_class.__name__}. "
+                f"Each concrete class must have a unique CONFIG_NAME."
+            )
+        UserExtendableNamedConfigMixin._registered_config_names[cls.CONFIG_NAME] = cls
+
+    @classmethod
+    def subclass_from_config_name(cls: type[_V], config_name: str) -> type[_V]:
+        """Create an instance of subclass by calling its string name (``CONFIG_NAME``).
+
+        This is and should stay dynamic because a user may interactively add subclasses
+        through REPL systems like Jupyter.
+
+        Args:
+            config_name: should match a subclass's ``CONFIG_NAME``.
+        """
+        subclass_config_name_map = {}
+        for type_ in _get_all_subclass_descendants(cls):
+            if name := getattr(type_, "CONFIG_NAME", ""):
+                subclass_config_name_map[name] = type_
+        try:
+            return subclass_config_name_map[config_name]
+        except KeyError as exc:
+            raise KeyError(
+                "Passed config_name does not match any subclass CONFIG_NAME"
+            ) from exc
+
+    @classmethod
+    def available_config_names(cls: type[_V]) -> list[str]:
+        """Get all available config names from the subclasses."""
+        config_names = [
+            type_.CONFIG_NAME
+            for type_ in _get_all_subclass_descendants(cls)
+            if hasattr(type_, "CONFIG_NAME") and type_.CONFIG_NAME
+        ]
+        return config_names

datacustomcode/run.py

@@ -0,0 +1,47 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import importlib
+import runpy
+from typing import List, Union
+
+from datacustomcode.config import config
+
+
+def run_entrypoint(
+    entrypoint: str, config_file: Union[str, None], dependencies: List[str]
+) -> None:
+    """Run the entrypoint script with the given config and dependencies.
+
+    Args:
+        entrypoint: The entrypoint script to run.
+        config_file: The config file to use.
+        dependencies: The dependencies to import.
+    """
+    if config_file:
+        config.load(config_file)
+    for dependency in dependencies:
+        try:
+            importlib.import_module(dependency)
+        except ModuleNotFoundError as exc:
+            try:
+                if "." in dependency:
+                    module_name, object_name = dependency.rsplit(".", 1)
+                    module = importlib.import_module(module_name)
+                    getattr(module, object_name)
+                else:
+                    raise exc
+            except (ModuleNotFoundError, AttributeError) as inner_exc:
+                raise inner_exc from exc
+    runpy.run_path(entrypoint, init_globals=globals(), run_name="__main__")

datacustomcode/scan.py

@@ -0,0 +1,153 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import ast
+from typing import (
+    Dict,
+    List,
+    Union,
+)
+
+import pydantic
+
+DATA_ACCESS_METHODS = ["read_dlo", "read_dmo", "write_to_dlo", "write_to_dmo"]
+
+
+class DataAccessLayerCalls(pydantic.BaseModel):
+    read_dlo: frozenset[str]
+    read_dmo: frozenset[str]
+    write_to_dlo: frozenset[str]
+    write_to_dmo: frozenset[str]
+
+    @pydantic.model_validator(mode="after")
+    def validate_access_layer(self) -> DataAccessLayerCalls:
+        if self.read_dlo and self.read_dmo:
+            raise ValueError("Cannot read from DLO and DMO in the same file.")
+        if len(self.write_to_dlo) > 1 or len(self.write_to_dmo) > 1:
+            raise ValueError(
+                "Cannot write to more than one DLO or DMO in the same file."
+            )
+        if not self.read_dlo and not self.read_dmo:
+            raise ValueError("Must read from at least one DLO or DMO.")
+        if self.read_dlo and self.write_to_dmo:
+            raise ValueError("Cannot read from DLO and write to DMO in the same file.")
+        if self.read_dmo and self.write_to_dlo:
+            raise ValueError("Cannot read from DMO and write to DLO in the same file.")
+        return self
+
+    @property
+    def input_str(self) -> str:
+        if self.read_dlo:
+            return next(iter(self.read_dlo))
+        return next(iter(self.read_dmo))
+
+    @property
+    def output_str(self) -> str:
+        if self.write_to_dlo:
+            return next(iter(self.write_to_dlo))
+        return next(iter(self.write_to_dmo))
+
+
+class ClientMethodVisitor(ast.NodeVisitor):
+    """AST Visitor that finds all instances of Client read/write method calls."""
+
+    def __init__(self) -> None:
+        self._read_dlo_instances: set[str] = set()
+        self._read_dmo_instances: set[str] = set()
+        self._write_to_dlo_instances: set[str] = set()
+        self._write_to_dmo_instances: set[str] = set()
+        self.variable_values: Dict[str, Union[str, None]] = {}
+
+    def visit_Assign(self, node: ast.Assign) -> None:
+        """Track variable assignments that might be used as DLO/DMO names."""
+        for target in node.targets:
+            if isinstance(target, ast.Name):
+                var_name = target.id
+                if isinstance(node.value, ast.Constant) and isinstance(
+                    node.value.value, str
+                ):
+                    self.variable_values[var_name] = node.value.value
+                else:
+                    self.variable_values[var_name] = None
+        self.generic_visit(node)
+
+    def visit_Call(self, node: ast.Call) -> None:
+        """Visit a method call and check if it's a Client read/write method."""
+        if isinstance(node.func, ast.Attribute) and isinstance(
+            node.func.value, ast.Name
+        ):
+            method_name = node.func.attr
+            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]
+
+                if name:
+                    if method_name == "read_dlo":
+                        self._read_dlo_instances.add(name)
+                    elif method_name == "read_dmo":
+                        self._read_dmo_instances.add(name)
+                    elif method_name == "write_to_dlo":
+                        self._write_to_dlo_instances.add(name)
+                    elif method_name == "write_to_dmo":
+                        self._write_to_dmo_instances.add(name)
+        self.generic_visit(node)
+
+    def found(self) -> DataAccessLayerCalls:
+        return DataAccessLayerCalls(
+            read_dlo=frozenset(self._read_dlo_instances),
+            read_dmo=frozenset(self._read_dmo_instances),
+            write_to_dlo=frozenset(self._write_to_dlo_instances),
+            write_to_dmo=frozenset(self._write_to_dmo_instances),
+        )
+
+
+def scan_file(file_path: str) -> DataAccessLayerCalls:
+    """Scan a single Python file for Client read/write method calls."""
+    with open(file_path, "r") as f:
+        code = f.read()
+        tree = ast.parse(code)
+        visitor = ClientMethodVisitor()
+        visitor.visit(tree)
+        return visitor.found()
+
+
+def dc_config_json_from_file(file_path: str) -> dict:
+    """Create a Data Cloud Custom Code config JSON from a script."""
+    output = scan_file(file_path)
+    read = {}
+    if output.read_dlo:
+        read["dlo"] = list(output.read_dlo)
+    else:
+        read["dmo"] = list(output.read_dmo)
+    write = {}
+    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,
+        },
+    }
+    return config

datacustomcode/template.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import shutil
+
+from loguru import logger
+
+template_dir = os.path.join(os.path.dirname(__file__), "templates")
+
+
+def copy_template(target_dir: str) -> None:
+    """Copy the template to the target directory."""
+    os.makedirs(target_dir, exist_ok=True)
+
+    for item in os.listdir(template_dir):
+        source = os.path.join(template_dir, item)
+        destination = os.path.join(target_dir, item)
+
+        if os.path.isdir(source):
+            logger.debug(f"Copying directory {source} to {destination}...")
+            shutil.copytree(source, destination, dirs_exist_ok=True)
+        else:
+            logger.debug(f"Copying file {source} to {destination}...")
+            shutil.copy2(source, destination)

datacustomcode/templates/.devcontainer/devcontainer.json

@@ -0,0 +1,10 @@
+{
+	"name": "Existing Dockerfile",
+	"build": {
+		"context": "..",
+		"dockerfile": "../Dockerfile"
+	},
+	"features": {
+		"ghcr.io/devcontainers/features/git:1": {},
+	}
+}

datacustomcode/templates/Dockerfile

@@ -0,0 +1,20 @@
+FROM public.ecr.aws/emr-on-eks/spark/emr-7.3.0:latest
+
+USER root
+
+ENV ENVIRONMENT=Outside_DataCloud
+
+# install from dev requirements.txt
+COPY requirements-dev.txt ./requirements-dev.txt
+RUN pip3 install --no-cache-dir -r requirements-dev.txt
+
+# Install from requirements.txt:
+COPY requirements.txt ./requirements.txt
+RUN pip3 install --no-cache-dir -r requirements.txt
+
+# Create workspace directory
+RUN mkdir /workspace
+
+# Set user and working directory
+USER hadoop:hadoop
+WORKDIR /workspace

datacustomcode/templates/jupyterlab.sh

@@ -0,0 +1,97 @@
+#!/bin/bash
+
+# Description: This script is used to start/stop the jupyter notebook in a docker container
+
+# Function to open browser based on OS
+open_browser() {
+    local url=$1
+    case "$(uname -s)" in
+        Darwin*)  # macOS
+            open "$url"
+            ;;
+        Linux*)   # Linux
+            if command -v xdg-open &> /dev/null; then
+                xdg-open "$url"
+            elif command -v gnome-open &> /dev/null; then
+                gnome-open "$url"
+            else
+                echo "Could not detect the web browser to use"
+                return 1
+            fi
+            ;;
+        CYGWIN*|MINGW32*|MSYS*|MINGW*)  # Windows
+            start "$url"
+            ;;
+        *)
+            echo "Unknown operating system"
+            return 1
+            ;;
+    esac
+}
+
+# Function to check if Docker is installed and running
+check_docker() {
+    if ! command -v docker &> /dev/null; then
+        echo "Docker is not installed"
+        exit 1
+    fi
+    echo "Docker is installed"
+    docker --version
+
+    if ! docker info &> /dev/null; then
+        echo "Docker daemon is not running"
+        exit 1
+    fi
+    echo "Docker daemon is running"
+}
+
+# Function to start Jupyter server
+start_jupyter() {
+    echo "Building the docker image"
+    docker build -t datacloud-byoc .
+
+    echo "Running the docker container"
+    docker run -d --rm -p 8888:8888 \
+        -v $(pwd):/workspace \
+        --name jupyter-server \
+        datacloud-byoc jupyter lab \
+        --ip=0.0.0.0 \
+        --port=8888 \
+        --no-browser \
+        --allow-root \
+        --NotebookApp.token='' \
+        --NotebookApp.password='' \
+        --notebook-dir=/workspace
+
+    sleep 3  # Wait for server to start
+    open_browser "http://localhost:8888"
+}
+
+# Function to stop Jupyter server
+stop_jupyter() {
+    echo "Stopping Jupyter server container..."
+    if docker ps -q --filter "name=jupyter-server" | grep -q .; then
+        docker stop jupyter-server
+        echo "Jupyter server stopped successfully"
+    else
+        echo "No Jupyter server container running"
+    fi
+}
+
+# Main script logic
+case "$1" in
+    "start")
+        check_docker
+        start_jupyter
+        ;;
+    "stop")
+        check_docker
+        stop_jupyter
+        ;;
+    *)
+        echo "Usage: $0 {start|stop}"
+        echo "  start - Start Jupyter server"
+        echo "  stop  - Stop Jupyter server"
+        exit 1
+        ;;
+esac

datacustomcode/templates/payload/config.json

@@ -0,0 +1 @@
+{}

datacustomcode/templates/payload/entrypoint.py

@@ -0,0 +1,10 @@
+from datacustomcode.client import Client
+
+
+def main():
+    client = Client()  # noqa: F841
+    # TODO: Add your custom code here
+
+
+if __name__ == "__main__":
+    main()

datacustomcode/templates/requirements-dev.txt

@@ -0,0 +1,10 @@
+# Required packages for the project - Do not modify
+salesforce-cdp-connector>=1.0.16
+pyspark
+pandas
+numpy
+pydantic
+jupyterlab
+ipywidgets
+tqdm
+salesforce-data-customcode

datacustomcode/templates/requirements.txt

@@ -0,0 +1 @@
+# Packages required for the custom code