salesforce-data-customcode 3.0.1__tar.gz → 4.0.0__tar.gz

{salesforce_data_customcode-3.0.1 → salesforce_data_customcode-4.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: salesforce-data-customcode
-Version: 3.0.1
+Version: 4.0.0
 Summary: Data Cloud Custom Code SDK
 License-Expression: Apache-2.0
 License-File: LICENSE.txt
@@ -18,6 +18,7 @@ Requires-Dist: pandas
 Requires-Dist: pydantic (==2.13.1)
 Requires-Dist: pyspark (==3.5.1)
 Requires-Dist: pyyaml (>=6.0,<7.0)
+Requires-Dist: requests (==2.33.1)
 Requires-Dist: salesforce-cdp-connector (>=1.0.19)
 Requires-Dist: setuptools_scm (>=7.1.0,<8.0.0)
 Description-Content-Type: text/markdown

{salesforce_data_customcode-3.0.1 → salesforce_data_customcode-4.0.0}/pyproject.toml

@@ -18,7 +18,7 @@ license = "Apache-2.0"
 name = "salesforce-data-customcode"
 readme = "README.md"
 requires-python = ">=3.10,<3.12"
-version = "3.0.1"
+version = "4.0.0"
 
 [tool.black]
 exclude = '''
@@ -105,6 +105,7 @@ pydantic = "2.13.1"
 pyspark = "3.5.1"
 python = ">=3.10,<3.12"
 pyyaml = "^6.0"
+requests = "2.33.1"
 salesforce-cdp-connector = ">=1.0.19"
 setuptools_scm = "^7.1.0"
 

{salesforce_data_customcode-3.0.1 → salesforce_data_customcode-4.0.0}/src/datacustomcode/__init__.py

@@ -17,15 +17,11 @@ from datacustomcode.client import Client
 from datacustomcode.credentials import AuthType, Credentials
 from datacustomcode.io.reader.query_api import QueryAPIDataCloudReader
 from datacustomcode.io.writer.print import PrintDataCloudWriter
-from datacustomcode.proxy.client.LocalProxyClientProvider import (
-    LocalProxyClientProvider,
-)
 
 __all__ = [
     "AuthType",
     "Client",
     "Credentials",
-    "LocalProxyClientProvider",
     "PrintDataCloudWriter",
     "QueryAPIDataCloudReader",
 ]

{salesforce_data_customcode-3.0.1 → salesforce_data_customcode-4.0.0}/src/datacustomcode/auth.py

@@ -170,7 +170,7 @@ def do_oauth_browser_flow(
 
     # Start callback server
     click.echo(f"\nStarting local callback server on {redirect_uri}...")
-    server, actual_port = _run_oauth_callback_server(redirect_uri, auth_code_queue)
+    server, _actual_port = _run_oauth_callback_server(redirect_uri, auth_code_queue)
 
     # Build authorization URL with final redirect_uri
     auth_url = (

{salesforce_data_customcode-3.0.1 → salesforce_data_customcode-4.0.0}/src/datacustomcode/cli.py

@@ -27,6 +27,13 @@ from loguru import logger
 
 from datacustomcode import AuthType
 from datacustomcode.auth import configure_oauth_tokens
+from datacustomcode.constants import (
+    CONFIG_FILE,
+    ENTRYPOINT_FILE,
+    PAYLOAD_DIR,
+    TEST_FILE,
+    TESTS_DIR,
+)
 from datacustomcode.scan import find_base_directory, get_package_type
 
 
@@ -74,6 +81,30 @@ def _configure_client_credentials(
     )
 
 
+def _generate_function_test_file(entrypoint_path: str) -> Optional[str]:
+    """Generate test.json file for a function.
+
+    Args:
+        entrypoint_path: Path to the function's entrypoint.py
+
+    Returns:
+        Path to generated test.json, or None if generation failed
+    """
+    from datacustomcode.function_utils import generate_test_json
+
+    tests_dir = os.path.join(os.path.dirname(entrypoint_path), TESTS_DIR)
+    os.makedirs(tests_dir, exist_ok=True)
+    test_json_path = os.path.join(tests_dir, TEST_FILE)
+
+    try:
+        generate_test_json(entrypoint_path, test_json_path)
+        logger.debug(f"Generated test JSON at {test_json_path}")
+        return test_json_path
+    except Exception as e:
+        logger.warning(f"Could not generate test.json: {e}")
+        return None
+
+
 @cli.command()
 @click.option("--profile", default="default", help="Credential profile name")
 @click.option(
@@ -162,7 +193,6 @@ def zip(path: str, network: str):
 
     Choose based on your workload requirements.""",
 )
-@click.option("--function-invoke-opt")
 @click.option(
     "--sf-cli-org",
     default=None,
@@ -176,16 +206,18 @@ def deploy(
     cpu_size: str,
     profile: str,
     network: str,
-    function_invoke_opt: str,
     sf_cli_org: Optional[str],
 ):
-    from datacustomcode.credentials import Credentials
+    from datacustomcode.constants import USE_IN_FEATURE_MAPPING_FOR_CONNECT_API
     from datacustomcode.deploy import (
         COMPUTE_TYPES,
-        AccessTokenResponse,
         CodeExtensionMetadata,
-        _retrieve_access_token_from_sf_cli,
         deploy_full,
+        infer_use_in_feature,
+    )
+    from datacustomcode.token_provider import (
+        CredentialsTokenProvider,
+        SFCLITokenProvider,
     )
 
     logger.debug("Deploying project")
@@ -210,32 +242,34 @@ def deploy(
     )
 
     if package_type == "function":
-        if not function_invoke_opt:
+        # Infer use_in_feature from function signature
+        entrypoint_path = os.path.join(path, ENTRYPOINT_FILE)
+        use_in_feature = infer_use_in_feature(entrypoint_path)
+        if use_in_feature:
+            logger.info(f"Inferred use_in_feature: {use_in_feature}")
+        else:
             click.secho(
-                "Error: Function invoke options are required for function package type",
+                "Error: Could not infer function invoke options. "
+                "Please provide --use-in-feature",
                 fg="red",
             )
             raise click.Abort()
+
+        # Map user-provided feature names to API names
+        mapped_feature = USE_IN_FEATURE_MAPPING_FOR_CONNECT_API.get(
+            use_in_feature, use_in_feature
+        )
+        metadata.functionInvokeOptions = [mapped_feature]
+
+    try:
+        if sf_cli_org:
+            auth = SFCLITokenProvider(sf_cli_org).get_token()
         else:
-            function_invoke_options = function_invoke_opt.split(",")
-            metadata.functionInvokeOptions = function_invoke_options
-
-    auth: Union[Credentials, AccessTokenResponse]
-    if sf_cli_org:
-        try:
-            auth = _retrieve_access_token_from_sf_cli(sf_cli_org)
-        except RuntimeError as e:
-            click.secho(f"Error: {e}", fg="red")
-            raise click.Abort() from None
-    else:
-        try:
-            auth = Credentials.from_available(profile=profile)
-        except ValueError as e:
-            click.secho(
-                f"Error: {e}",
-                fg="red",
-            )
-            raise click.Abort() from None
+            auth = CredentialsTokenProvider(profile).get_token()
+    except RuntimeError as e:
+        click.secho(f"Error: {e}", fg="red")
+        raise click.Abort() from None
+
     deploy_full(path, metadata, auth, network)
 
 
@@ -244,7 +278,12 @@ def deploy(
 @click.option(
     "--code-type", default="script", type=click.Choice(["script", "function"])
 )
-def init(directory: str, code_type: str):
+@click.option(
+    "--use-in-feature",
+    default="SearchIndexChunking",
+    help="Feature where this function will be used (only applicable for function).",
+)
+def init(directory: str, code_type: str, use_in_feature: Optional[str]):
     from datacustomcode.scan import (
         dc_config_json_from_file,
         update_config,
@@ -256,9 +295,9 @@ def init(directory: str, code_type: str):
     if code_type == "script":
         copy_script_template(directory)
     elif code_type == "function":
-        copy_function_template(directory)
-    entrypoint_path = os.path.join(directory, "payload", "entrypoint.py")
-    config_location = os.path.join(os.path.dirname(entrypoint_path), "config.json")
+        copy_function_template(directory, use_in_feature)
+    entrypoint_path = os.path.join(directory, PAYLOAD_DIR, ENTRYPOINT_FILE)
+    config_location = os.path.join(os.path.dirname(entrypoint_path), CONFIG_FILE)
 
     # Write package type to SDK-specific config
     sdk_config = {"type": code_type}
@@ -271,6 +310,7 @@ def init(directory: str, code_type: str):
     updated_config_json = update_config(entrypoint_path)
     with open(config_location, "w") as f:
         json.dump(updated_config_json, f, indent=2)
+
     click.echo(
         "Start developing by updating the code in "
         + click.style(entrypoint_path, fg="blue", bold=True)
@@ -281,6 +321,24 @@ def init(directory: str, code_type: str):
         + " to automatically update config.json when you make changes to your code"
     )
 
+    # Generate test.json for functions
+    if code_type == "function":
+        test_json_path = _generate_function_test_file(entrypoint_path)
+        if test_json_path:
+            click.echo(
+                "Generated test file at "
+                + click.style(test_json_path, fg="blue", bold=True)
+            )
+            click.echo(
+                "Test your function locally with "
+                + click.style(
+                    f"datacustomcode run {entrypoint_path} "
+                    f"--test-with {test_json_path}",
+                    fg="blue",
+                    bold=True,
+                )
+            )
+
 
 @cli.command()
 @click.argument("filename")
@@ -292,7 +350,7 @@ def init(directory: str, code_type: str):
 def scan(filename: str, config: str, dry_run: bool, no_requirements: bool):
     from datacustomcode.scan import update_config, write_requirements_file
 
-    config_location = config or os.path.join(os.path.dirname(filename), "config.json")
+    config_location = config or os.path.join(os.path.dirname(filename), CONFIG_FILE)
     click.echo(
         "Dumping scan results to config file: "
         + click.style(config_location, fg="blue", bold=True)
@@ -318,6 +376,12 @@ def scan(filename: str, config: str, dry_run: bool, no_requirements: bool):
 @click.option("--config-file", default=None)
 @click.option("--dependencies", default=[], multiple=True)
 @click.option("--profile", default="default")
+@click.option(
+    "--test-with",
+    default=None,
+    type=click.Path(exists=True),
+    help="Path to test JSON file for function testing",
+)
 @click.option(
     "--sf-cli-org",
     default=None,
@@ -328,10 +392,16 @@ def run(
     config_file: Union[str, None],
     dependencies: List[str],
     profile: str,
+    test_with: Optional[str],
     sf_cli_org: Optional[str],
 ):
     from datacustomcode.run import run_entrypoint
 
     run_entrypoint(
-        entrypoint, config_file, dependencies, profile, sf_cli_org=sf_cli_org
+        entrypoint,
+        config_file,
+        dependencies,
+        profile,
+        test_file=test_with,
+        sf_cli_org=sf_cli_org,
     )

{salesforce_data_customcode-3.0.1 → salesforce_data_customcode-4.0.0}/src/datacustomcode/client.py

@@ -33,7 +33,6 @@ if TYPE_CHECKING:
 
     from datacustomcode.io.reader.base import BaseDataCloudReader
     from datacustomcode.io.writer.base import BaseDataCloudWriter, WriteMode
-    from datacustomcode.proxy.client.base import BaseProxyClient
     from datacustomcode.spark.base import BaseSparkSessionProvider
 
 
@@ -107,7 +106,6 @@ class Client:
     _reader: BaseDataCloudReader
     _writer: BaseDataCloudWriter
     _file: DefaultFindFilePath
-    _proxy: Optional[BaseProxyClient]
     _data_layer_history: dict[DataCloudObjectType, set[str]]
     _code_type: str
 
@@ -115,7 +113,6 @@ class Client:
         cls,
         reader: Optional[BaseDataCloudReader] = None,
         writer: Optional["BaseDataCloudWriter"] = None,
-        proxy: Optional[BaseProxyClient] = None,
         spark_provider: Optional["BaseSparkSessionProvider"] = None,
         code_type: str = "script",
     ) -> Client:
@@ -223,11 +220,6 @@ class Client:
         self._validate_data_layer_history_does_not_contain(DataCloudObjectType.DLO)
         return self._writer.write_to_dmo(name, dataframe, write_mode, **kwargs)  # type: ignore[no-any-return]
 
-    def call_llm_gateway(self, LLM_MODEL_ID: str, prompt: str, maxTokens: int) -> str:
-        if self._proxy is None:
-            raise ValueError("No proxy configured; set proxy or proxy_config")
-        return self._proxy.call_llm_gateway(LLM_MODEL_ID, prompt, maxTokens)  # type: ignore[no-any-return]
-
     def find_file_path(self, file_name: str) -> Path:
         """Return a file path"""
 

{salesforce_data_customcode-3.0.1 → salesforce_data_customcode-4.0.0}/src/datacustomcode/cmd.py

@@ -104,6 +104,6 @@ def _cmd_output(
 
 
 def cmd_output(*cmd: str, **kwargs: Any) -> Union[str, None]:
-    returncode, stdout_b, stderr_b = _cmd_output(*cmd, **kwargs)
+    _returncode, stdout_b, _stderr_b = _cmd_output(*cmd, **kwargs)
     stdout = stdout_b.decode() if stdout_b is not None else None
     return stdout

salesforce_data_customcode-4.0.0/src/datacustomcode/common_config.py

@@ -0,0 +1,63 @@
+# 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 abc import ABC, abstractmethod
+import os
+from typing import Any
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+)
+import yaml
+
+DEFAULT_CONFIG_NAME = "config.yaml"
+
+
+def default_config_file() -> str:
+    return os.path.join(os.path.dirname(__file__), DEFAULT_CONFIG_NAME)
+
+
+class ForceableConfig(BaseModel):
+    force: bool = Field(
+        default=False,
+        description="If True, this takes precedence over parameters passed to the "
+        "initializer of the client",
+    )
+
+
+class BaseObjectConfig(ForceableConfig):
+    model_config = ConfigDict(validate_default=True, extra="forbid")
+    type_config_name: str = Field(
+        description="The config name of the object to create",
+    )
+    options: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Options passed to the constructor.",
+    )
+
+
+class BaseConfig(ABC, BaseModel):
+    @abstractmethod
+    def update(self, other: Any) -> "BaseConfig": ...
+
+    def load(self, config_path: str) -> "BaseConfig":
+        """Load configuration from a YAML file and merge with existing config"""
+        with open(config_path, "r") as f:
+            config_data = yaml.safe_load(f)
+
+        loaded_config = self.__class__.model_validate(config_data)
+        self.update(loaded_config)
+        return self

{salesforce_data_customcode-3.0.1 → salesforce_data_customcode-4.0.0}/src/datacustomcode/config.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 from __future__ import annotations
 
-import os
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -26,56 +25,36 @@ from typing import (
     cast,
 )
 
-from pydantic import (
-    BaseModel,
-    ConfigDict,
-    Field,
+from pydantic import Field
+
+from datacustomcode.common_config import (
+    BaseConfig,
+    BaseObjectConfig,
+    ForceableConfig,
+    default_config_file,
 )
-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: TCH002
-from datacustomcode.io.writer.base import BaseDataCloudWriter  # noqa: TCH002
-from datacustomcode.proxy.base import BaseProxyAccessLayer
-from datacustomcode.proxy.client.base import BaseProxyClient  # noqa: TCH002
 from datacustomcode.spark.base import BaseSparkSessionProvider
 
-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.",
-    )
+    from datacustomcode.io.reader.base import BaseDataCloudReader
+    from datacustomcode.io.writer.base import BaseDataCloudWriter
 
 
 _T = TypeVar("_T", bound="BaseDataAccessLayer")
 
 
-class AccessLayerObjectConfig(ForceableConfig, Generic[_T]):
-    model_config = ConfigDict(validate_default=True, extra="forbid")
+class AccessLayerObjectConfig(BaseObjectConfig, Generic[_T]):
     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))
+        return cast("_T", type_(spark=spark, **self.options))
 
 
 class SparkConfig(ForceableConfig):
@@ -94,41 +73,18 @@ class SparkConfig(ForceableConfig):
 
 _P = TypeVar("_P", bound=BaseSparkSessionProvider)
 
-_PX = TypeVar("_PX", bound=BaseProxyAccessLayer)
-
-
-class ProxyAccessLayerObjectConfig(ForceableConfig, Generic[_PX]):
-    """Config for proxy clients that take no constructor args (e.g. no spark)."""
-
-    model_config = ConfigDict(validate_default=True, extra="forbid")
-    type_base: ClassVar[Type[BaseProxyAccessLayer]] = BaseProxyAccessLayer
-    type_config_name: str = Field(
-        description="CONFIG_NAME of the proxy client (e.g. 'LocalProxyClient').",
-    )
-    options: dict[str, Any] = Field(default_factory=dict)
-
-    def to_object(self) -> _PX:
-        type_ = self.type_base.subclass_from_config_name(self.type_config_name)
-        return cast(_PX, type_(**self.options))
-
 
-class SparkProviderConfig(ForceableConfig, Generic[_P]):
-    model_config = ConfigDict(validate_default=True, extra="forbid")
+class SparkProviderConfig(BaseObjectConfig, Generic[_P]):
     type_base: ClassVar[Type[BaseSparkSessionProvider]] = BaseSparkSessionProvider
-    type_config_name: str = Field(
-        description="CONFIG_NAME of the Spark session provider."
-    )
-    options: dict[str, Any] = Field(default_factory=dict)
 
     def to_object(self) -> _P:
         type_ = self.type_base.subclass_from_config_name(self.type_config_name)
-        return cast(_P, type_(**self.options))
+        return cast("_P", type_(**self.options))
 
 
-class ClientConfig(BaseModel):
-    reader_config: Union[AccessLayerObjectConfig[BaseDataCloudReader], None] = None
-    writer_config: Union[AccessLayerObjectConfig[BaseDataCloudWriter], None] = None
-    proxy_config: Union[ProxyAccessLayerObjectConfig[BaseProxyClient], None] = None
+class ClientConfig(BaseConfig):
+    reader_config: Union[AccessLayerObjectConfig["BaseDataCloudReader"], None] = None
+    writer_config: Union[AccessLayerObjectConfig["BaseDataCloudWriter"], None] = None
     spark_config: Union[SparkConfig, None] = None
     spark_provider_config: Union[
         SparkProviderConfig[BaseSparkSessionProvider], None
@@ -156,38 +112,16 @@ class ClientConfig(BaseModel):
 
         self.reader_config = merge(self.reader_config, other.reader_config)
         self.writer_config = merge(self.writer_config, other.writer_config)
-        self.proxy_config = merge(self.proxy_config, other.proxy_config)
         self.spark_config = merge(self.spark_config, other.spark_config)
         self.spark_provider_config = merge(
             self.spark_provider_config, other.spark_provider_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())
+config = ClientConfig()
+config.load(default_config_file())

{salesforce_data_customcode-3.0.1 → salesforce_data_customcode-4.0.0}/src/datacustomcode/config.yaml

@@ -19,7 +19,12 @@ spark_config:
     spark.sql.execution.arrow.pyspark.enabled: 'true'
     spark.driver.extraJavaOptions: -Djava.security.manager=allow
 
-proxy_config:
-  type_config_name: LocalProxyClientProvider
+einstein_predictions_config:
+  type_config_name: DefaultEinsteinPredictions
+  options:
+    credentials_profile: default
+
+llm_gateway_config:
+  type_config_name: DefaultLLMGateway
   options:
     credentials_profile: default

salesforce_data_customcode-4.0.0/src/datacustomcode/constants.py

@@ -0,0 +1,45 @@
+# 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.
+
+"""Constants used throughout the datacustomcode package."""
+
+# File and directory names
+ENTRYPOINT_FILE = "entrypoint.py"
+CONFIG_FILE = "config.json"
+PAYLOAD_DIR = "payload"
+TESTS_DIR = "tests"
+TEST_FILE = "test.json"
+REQUIREMENTS_FILE = "requirements.txt"
+
+# Default values
+DEFAULT_PROFILE = "default"
+DEFAULT_NETWORK = "default"
+DEFAULT_CPU_SIZE = "CPU_2XL"
+
+# Feature to template folder mapping
+FEATURE_TEMPLATE_MAPPING = {
+    "SearchIndexChunking": "chunking",
+}
+
+# Feature name to Connect API name mapping
+USE_IN_FEATURE_MAPPING_FOR_CONNECT_API = {
+    "SearchIndexChunking": "UnstructuredChunking",
+}
+
+# Pydantic request/response type names to feature names
+REQUEST_TYPE_TO_FEATURE = {
+    "SearchIndexChunkingV1Request": "SearchIndexChunking",
+    "SearchIndexChunkingV1Response": "SearchIndexChunking",
+}