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

datacustomcode/__init__.py

@@ -0,0 +1,20 @@
+# 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 datacustomcode.client import Client
+from datacustomcode.io.reader.query_api import QueryAPIDataCloudReader
+from datacustomcode.io.writer.print import PrintDataCloudWriter
+
+__all__ = ["Client", "QueryAPIDataCloudReader", "PrintDataCloudWriter"]

datacustomcode/cli.py

@@ -0,0 +1,142 @@
+# 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 importlib import metadata
+import json
+import os
+import sys
+from typing import List, Union
+
+import click
+from loguru import logger
+
+
+@click.group()
+@click.option("--debug", is_flag=True)
+def cli(debug: bool):
+    logger.remove()
+    if debug:
+        logger.configure(handlers=[{"sink": sys.stderr, "level": "DEBUG"}])
+    else:
+        logger.configure(handlers=[{"sink": sys.stderr, "level": "INFO"}])
+
+
+@cli.command()
+def version():
+    """Display the current version of the package."""
+    print(__name__)
+    try:
+        version = metadata.version("salesforce-data-customcode")
+        click.echo(f"salesforce-data-customcode version: {version}")
+    except metadata.PackageNotFoundError:
+        click.echo("Version information not available")
+
+
+@cli.command()
+@click.option("--profile", default="default")
+@click.option("--username", prompt=True)
+@click.option("--password", prompt=True, hide_input=True)
+@click.option("--client-id", prompt=True)
+@click.option("--client-secret", prompt=True)
+@click.option("--login-url", prompt=True)
+def configure(
+    username: str,
+    password: str,
+    client_id: str,
+    client_secret: str,
+    login_url: str,
+    profile: str,
+) -> None:
+    from datacustomcode.credentials import Credentials
+
+    Credentials(
+        username=username,
+        password=password,
+        client_id=client_id,
+        client_secret=client_secret,
+        login_url=login_url,
+    ).update_ini(profile=profile)
+
+
+@cli.command()
+@click.option("--profile", default="default")
+@click.option("--path", default="payload")
+@click.option("--name", required=True)
+@click.option("--version", default="0.0.1")
+@click.option("--description", default="Custom Data Transform Code")
+def deploy(profile: str, path: str, name: str, version: str, description: str):
+    from datacustomcode.credentials import Credentials
+    from datacustomcode.deploy import TransformationJobMetadata, deploy_full
+
+    logger.debug("Deploying project")
+
+    metadata = TransformationJobMetadata(
+        name=name,
+        version=version,
+        description=description,
+    )
+    try:
+        credentials = Credentials.from_ini(profile=profile)
+    except KeyError:
+        click.secho(
+            f"Error: Profile {profile} not found in credentials.ini. "
+            "Run `datacustomcode configure` to create a credentialsprofile.",
+            fg="red",
+        )
+        raise click.Abort() from None
+    deploy_full(path, metadata, credentials)
+
+
+@cli.command()
+@click.argument("directory", default=".")
+def init(directory: str):
+    from datacustomcode.template import copy_template
+
+    click.echo("Copying template to " + click.style(directory, fg="blue", bold=True))
+    copy_template(directory)
+    click.echo(
+        "Start developing by updating the code in "
+        + click.style(f"{directory}/payload/entrypoint.py", fg="blue", bold=True)
+    )
+
+
+@cli.command()
+@click.argument("filename")
+@click.option("--config")
+@click.option("--dry-run", is_flag=True)
+def scan(filename: str, config: str, dry_run: bool):
+    from datacustomcode.scan import dc_config_json_from_file
+
+    config_location = config or os.path.join(os.path.dirname(filename), "config.json")
+    click.echo(
+        "Dumping scan results to config file: "
+        + click.style(config_location, fg="blue", bold=True)
+    )
+    click.echo("Scanning " + click.style(filename, fg="blue", bold=True) + "...")
+    config_json = dc_config_json_from_file(filename)
+
+    click.secho(json.dumps(config_json, indent=2), fg="yellow")
+    if not dry_run:
+        with open(config_location, "w") as f:
+            json.dump(config_json, f, indent=2)
+
+
+@cli.command()
+@click.argument("entrypoint")
+@click.option("--config-file", default=None)
+@click.option("--dependencies", default=[], multiple=True)
+def run(entrypoint: str, config_file: Union[str, None], dependencies: List[str]):
+    from datacustomcode.run import run_entrypoint
+
+    run_entrypoint(entrypoint, config_file, dependencies)

datacustomcode/client.py

@@ -0,0 +1,227 @@
+# 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 enum import Enum
+from typing import (
+    TYPE_CHECKING,
+    ClassVar,
+    Optional,
+)
+
+from pyspark.sql import SparkSession
+
+from datacustomcode.config import SparkConfig, config
+from datacustomcode.io.reader.base import BaseDataCloudReader
+
+if TYPE_CHECKING:
+    from pyspark.sql import DataFrame as PySparkDataFrame
+
+    from datacustomcode.io.reader.base import BaseDataCloudReader
+    from datacustomcode.io.writer.base import BaseDataCloudWriter, WriteMode
+
+
+def _setup_spark(spark_config: SparkConfig) -> SparkSession:
+    """Setup Spark session from config."""
+    builder = SparkSession.builder
+    if spark_config.master is not None:
+        builder = builder.master(spark_config.master)
+
+    builder = builder.appName(spark_config.app_name)
+    for key, value in spark_config.options.items():
+        builder = builder.config(key, value)
+    return builder.getOrCreate()
+
+
+class DataCloudObjectType(Enum):
+    DLO = "dlo"
+    DMO = "dmo"
+
+
+class DataCloudAccessLayerException(Exception):
+    """Exception raised when mixing DMOs and DLOs is detected."""
+
+    def __init__(
+        self,
+        data_layer_history: dict[DataCloudObjectType, set[str]],
+        should_not_contain: DataCloudObjectType,
+    ) -> None:
+        self.data_layer_history = data_layer_history
+        self.should_not_contain = should_not_contain
+
+    def __str__(self) -> str:
+        msg = (
+            "Mixed use of DMOs and DLOs. "
+            "You can only read from DMOs to write to DMOs "
+            "and read from DLOs to write to DLOs. "
+        )
+        if self.should_not_contain is DataCloudObjectType.DLO:
+            msg += (
+                "You have read from the following DLOs: "
+                f"{self.data_layer_history[DataCloudObjectType.DLO]} "
+                f"and are attempting to write to DMO. "
+            )
+        else:
+            msg += (
+                "You have read from the following DMOs: "
+                f"{self.data_layer_history[DataCloudObjectType.DMO]} "
+                f"and are attempting to write to to a DLO. "
+            )
+        msg += "Restart to clear history."
+        return msg
+
+
+class Client:
+    """Entrypoint for accessing DataCloud objects.
+
+    This is the object used to access Data Cloud DLOs and DMOs. Accessing DLOs/DMOs
+    are tracked and will throw an exception if they are mixed. In other words, you
+    can read from DLOs and write to DLOs, read from DMOs and write to DMOs, but you
+    cannot read from DLOs and write to DMOs or read from DMOs and write to DLOs.
+    Furthermore you cannot mix during merging tables. This class is a singleton to
+    prevent accidental mixing of DLOs and DMOs.
+
+    You can provide custom readers and writers to the client for advanced use
+    cases, but this is not recommended for testing as they may result in unexpected
+    behavior once deployed to Data Cloud. By default, the client intercepts all
+    read/write operations and mocks access to Data Cloud. For example, during
+    writing, we print to the console instead of writing to Data Cloud.
+
+    Args:
+        reader: A custom reader to use for reading Data Cloud objects.
+        writer: A custom writer to use for writing Data Cloud objects.
+
+    Example:
+    >>> client = Client()
+    >>> dlo = client.read_dlo("my_dlo")
+    >>> client.write_to_dmo("my_dmo", dlo)
+    """
+
+    _instance: ClassVar[Optional[Client]] = None
+    _reader: BaseDataCloudReader
+    _writer: BaseDataCloudWriter
+    _data_layer_history: dict[DataCloudObjectType, set[str]]
+
+    def __new__(
+        cls,
+        reader: Optional[BaseDataCloudReader] = None,
+        writer: Optional[BaseDataCloudWriter] = None,
+    ) -> Client:
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+
+            # Initialize Readers and Writers from config
+            # and/or provided reader and writer
+            if reader is None or writer is None:
+                # We need a spark because we will initialize readers and writers
+                if config.spark_config is None:
+                    raise ValueError(
+                        "Spark config is required when reader/writer is not provided"
+                    )
+                spark = _setup_spark(config.spark_config)
+
+            if config.reader_config is None and reader is None:
+                raise ValueError(
+                    "Reader config is required when reader is not provided"
+                )
+            elif reader is None or (
+                config.reader_config is not None and config.reader_config.force
+            ):
+                reader_init = config.reader_config.to_object(spark)  # type: ignore
+            else:
+                reader_init = reader
+            if config.writer_config is None and writer is None:
+                raise ValueError(
+                    "Writer config is required when writer is not provided"
+                )
+            elif writer is None or (
+                config.writer_config is not None and config.writer_config.force
+            ):
+                writer_init = config.writer_config.to_object(spark)  # type: ignore
+            else:
+                writer_init = writer
+            cls._instance._reader = reader_init
+            cls._instance._writer = writer_init
+            cls._instance._data_layer_history = {
+                DataCloudObjectType.DLO: set(),
+                DataCloudObjectType.DMO: set(),
+            }
+        elif (reader is not None or writer is not None) and cls._instance is not None:
+            raise ValueError("Cannot set reader or writer after client is initialized")
+        return cls._instance
+
+    def read_dlo(self, name: str) -> PySparkDataFrame:
+        """Read a DLO from Data Cloud.
+
+        Args:
+            name: The name of the DLO to read.
+
+        Returns:
+            A PySpark DataFrame containing the DLO data.
+        """
+        self._record_dlo_access(name)
+        return self._reader.read_dlo(name)
+
+    def read_dmo(self, name: str) -> PySparkDataFrame:
+        """Read a DMO from Data Cloud.
+
+        Args:
+            name: The name of the DMO to read.
+
+        Returns:
+            A PySpark DataFrame containing the DMO data.
+        """
+        self._record_dmo_access(name)
+        return self._reader.read_dmo(name)
+
+    def write_to_dlo(
+        self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs
+    ) -> None:
+        """Write a PySpark DataFrame to a DLO in Data Cloud.
+
+        Args:
+            name: The name of the DLO to write to.
+            dataframe: The PySpark DataFrame to write.
+            write_mode: The write mode to use for writing to the DLO.
+        """
+        self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DMO)
+        return self._writer.write_to_dlo(name, dataframe, write_mode, **kwargs)
+
+    def write_to_dmo(
+        self, name: str, dataframe: PySparkDataFrame, write_mode: WriteMode, **kwargs
+    ) -> None:
+        """Write a PySpark DataFrame to a DMO in Data Cloud.
+
+        Args:
+            name: The name of the DMO to write to.
+            dataframe: The PySpark DataFrame to write.
+            write_mode: The write mode to use for writing to the DMO.
+        """
+        self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DLO)
+        return self._writer.write_to_dmo(name, dataframe, write_mode, **kwargs)
+
+    def _validate_data_layer_history_does_not_contain(
+        self, data_cloud_object_type: DataCloudObjectType
+    ) -> None:
+        if len(self._data_layer_history[data_cloud_object_type]) > 0:
+            raise DataCloudAccessLayerException(
+                self._data_layer_history, data_cloud_object_type
+            )
+
+    def _record_dlo_access(self, name: str) -> None:
+        self._data_layer_history[DataCloudObjectType.DLO].add(name)
+
+    def _record_dmo_access(self, name: str) -> None:
+        self._data_layer_history[DataCloudObjectType.DMO].add(name)

datacustomcode/cmd.py

@@ -0,0 +1,105 @@
+# 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.
+"""
+This module is shamelessly copied from conda to nicely wrap subprocess calls.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import subprocess
+from typing import Any, Union
+
+
+def _force_bytes(exc: Any) -> bytes:
+    with contextlib.suppress(TypeError):
+        return bytes(exc)
+    with contextlib.suppress(Exception):
+        return str(exc).encode()
+    return f"<unprintable {type(exc).__name__} object>".encode()
+
+
+def _setdefault_kwargs(kwargs: dict[str, Any]) -> None:
+    for arg in ("stdin", "stdout", "stderr"):
+        kwargs.setdefault(arg, subprocess.PIPE)
+
+
+def _oserror_to_output(e: OSError) -> tuple[int, bytes, None]:
+    return 1, _force_bytes(e).rstrip(b"\n") + b"\n", None
+
+
+class CalledProcessError(RuntimeError):
+    """Nicely formatted subprocess call error."""
+
+    def __init__(
+        self,
+        returncode: int,
+        cmd: tuple[str, ...],
+        stdout: bytes,
+        stderr: Union[bytes, None],
+    ) -> None:
+        super().__init__(returncode, cmd, stdout, stderr)
+        self.returncode = returncode
+        self.cmd = cmd
+        self.stdout = stdout
+        self.stderr = stderr
+
+    def __bytes__(self) -> bytes:
+        def _indent_or_none(part: Union[bytes, None]) -> bytes:
+            if part:
+                return b"\n    " + part.replace(b"\n", b"\n    ").rstrip()
+            else:
+                return b" (none)"
+
+        return b"".join(
+            (
+                f"command: {self.cmd!r}\n".encode(),
+                f"return code: {self.returncode}\n".encode(),
+                b"stdout:",
+                self.stdout,
+                b"\n",
+                b"stderr:",
+                _indent_or_none(self.stderr),
+            )
+        )
+
+    def __str__(self) -> str:
+        return self.__bytes__().decode()
+
+
+def _cmd_output(
+    *cmd: str,
+    check: bool = True,
+    **kwargs: Any,
+) -> tuple[int, bytes, Union[bytes, None]]:
+    _setdefault_kwargs(kwargs)
+    try:
+        kwargs.setdefault("shell", True)
+        proc = subprocess.Popen(cmd, **kwargs)
+    except OSError as e:
+        returncode, stdout_b, stderr_b = _oserror_to_output(e)
+    else:
+        stdout_b, stderr_b = proc.communicate()
+        returncode = proc.returncode
+    if check and returncode:
+        raise CalledProcessError(returncode, cmd, stdout_b, stderr_b)
+
+    return returncode, stdout_b, stderr_b
+
+
+def cmd_output(*cmd: str, **kwargs: Any) -> Union[str, None]:
+    returncode, stdout_b, stderr_b = _cmd_output(*cmd, **kwargs)
+    stdout = stdout_b.decode() if stdout_b is not None else None
+    return stdout

datacustomcode/config.py

@@ -0,0 +1,149 @@
+# 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 os
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Generic,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+)
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+)
+import yaml
+
+# This lets all readers and writers to be findable via config
+from datacustomcode.io import *  # noqa: F403
+from datacustomcode.io.base import BaseDataAccessLayer
+from datacustomcode.io.reader.base import BaseDataCloudReader  # noqa: TCH001
+from datacustomcode.io.writer.base import BaseDataCloudWriter  # noqa: TCH001
+
+DEFAULT_CONFIG_NAME = "config.yaml"
+
+
+if TYPE_CHECKING:
+    from pyspark.sql import SparkSession
+
+
+class ForceableConfig(BaseModel):
+    force: bool = Field(
+        default=False,
+        description="If True, this takes precedence over parameters passed to the "
+        "initializer of the client.",
+    )
+
+
+_T = TypeVar("_T", bound="BaseDataAccessLayer")
+
+
+class AccessLayerObjectConfig(ForceableConfig, Generic[_T]):
+    model_config = ConfigDict(validate_default=True, extra="forbid")
+    type_base: ClassVar[Type[BaseDataAccessLayer]] = BaseDataAccessLayer
+    type_config_name: str = Field(
+        description="The config name of the object to create. "
+        "For metrics, this would might be 'ipmnormal'. For custom classes, you can "
+        "assign a name to a class variable `CONFIG_NAME` and reference it here.",
+    )
+    options: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Options passed to the constructor.",
+    )
+
+    def to_object(self, spark: SparkSession) -> _T:
+        type_ = self.type_base.subclass_from_config_name(self.type_config_name)
+        return cast(_T, type_(spark=spark, **self.options))
+
+
+class SparkConfig(ForceableConfig):
+    app_name: str = Field(
+        description="The name of the Spark application.",
+    )
+    master: Union[str, None] = Field(
+        default=None,
+        description="The Spark master URL.",
+    )
+    options: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Options passed to the SparkSession constructor.",
+    )
+
+
+class ClientConfig(BaseModel):
+    reader_config: Union[AccessLayerObjectConfig[BaseDataCloudReader], None] = None
+    writer_config: Union[AccessLayerObjectConfig[BaseDataCloudWriter], None] = None
+    spark_config: Union[SparkConfig, None] = None
+
+    def update(self, other: ClientConfig) -> ClientConfig:
+        """Merge this ClientConfig with another, respecting force flags.
+
+        Args:
+            other: Another ClientConfig to merge with this one
+
+        Returns:
+            Self, with updated values from the other config based on force flags.
+        """
+        TypeVarT = TypeVar("TypeVarT", bound=ForceableConfig)
+
+        def merge(
+            config_a: Union[TypeVarT, None], config_b: Union[TypeVarT, None]
+        ) -> Union[TypeVarT, None]:
+            if config_a is not None and config_a.force:
+                return config_a
+            if config_b:
+                return config_b
+            return config_a
+
+        self.reader_config = merge(self.reader_config, other.reader_config)
+        self.writer_config = merge(self.writer_config, other.writer_config)
+        self.spark_config = merge(self.spark_config, other.spark_config)
+        return self
+
+    def load(self, config_path: str) -> ClientConfig:
+        """Load a config from a file and update this config with it.
+
+        Args:
+            config_path: The path to the config file
+
+        Returns:
+            Self, with updated values from the loaded config.
+        """
+        with open(config_path, "r") as f:
+            config_data = yaml.safe_load(f)
+        loaded_config = ClientConfig.model_validate(config_data)
+
+        return self.update(loaded_config)
+
+
+config = ClientConfig()
+"""Global config object.
+
+This is the object that makes config accessible globally and globally mutable.
+"""
+
+
+def _defaults() -> str:
+    return os.path.join(os.path.dirname(__file__), DEFAULT_CONFIG_NAME)
+
+
+config.load(_defaults())

datacustomcode/config.yaml

@@ -0,0 +1,15 @@
+reader_config:
+  type_config_name: QueryAPIDataCloudReader
+
+writer_config:
+  type_config_name: PrintDataCloudWriter
+
+spark_config:
+  app_name: DC Custom Code Python SDK Testing
+  master: local[*]
+  options:
+    spark.driver.host: localhost
+    spark.driver.bindAddress: 127.0.0.1
+    spark.submit.deployMode: client
+    spark.sql.execution.arrow.pyspark.enabled: 'true'
+    spark.driver.extraJavaOptions: -Djava.security.manager=allow