Merge pull request #141 from diksha-sf/main

Add support for external-callout

Author: diksha-sf

src/datacustomcode/client.py

@@ -32,6 +32,7 @@
 from datacustomcode.file.path.default import DefaultFindFilePath
 from datacustomcode.io.reader.base import BaseDataCloudReader
 from datacustomcode.llm_gateway_config import spark_llm_gateway_config
+from datacustomcode.named_credential_config import spark_named_credential_config
 from datacustomcode.spark.default import DefaultSparkSessionProvider
 
 if TYPE_CHECKING:
@@ -49,6 +50,9 @@
     from datacustomcode.io.reader.base import BaseDataCloudReader
     from datacustomcode.io.writer.base import BaseDataCloudWriter, WriteMode
     from datacustomcode.llm_gateway.spark_base import SparkLLMGateway
+    from datacustomcode.named_credential.spark_base import SparkNamedCredential
+    from datacustomcode.named_credential.types.http_request import HTTPRequest
+    from datacustomcode.named_credential.types.http_response import HTTPResponse
     from datacustomcode.spark.base import BaseSparkSessionProvider
 
 
@@ -160,6 +164,21 @@ def _build_spark_einstein_predictions() -> "SparkEinsteinPredictions":
     return cfg.to_object()
 
 
+def _build_spark_named_credential() -> "SparkNamedCredential":
+    """Instantiate the SDK-configured :class:`SparkNamedCredential`.
+
+    Raises:
+        RuntimeError: If no ``spark_named_credential_config`` has been loaded.
+    """
+    cfg = spark_named_credential_config.spark_named_credential_config
+    if cfg is None:
+        raise RuntimeError(
+            "spark_named_credential_config is not configured. Add a "
+            "'spark_named_credential_config' section to config.yaml."
+        )
+    return cfg.to_object()
+
+
 def einstein_predict_col(
     model_api_name: str,
     prediction_type: "PredictionType",
@@ -209,6 +228,40 @@ def einstein_predict_col(
     )
 
 
+def named_credential_request_col(
+    request: "HTTPRequest",
+    body: Optional["Column"] = None,
+) -> "Column":
+    """Build a Spark Column that makes one Named Credential callout per row.
+
+    The endpoint, method, and headers are fixed for the call (taken from
+    ``request``); only ``body`` varies per row. Use this instead of
+    :meth:`Client.named_credential_request` when the callout runs across a
+    DataFrame so each row is dispatched independently rather than one-shot on
+    the driver.
+
+    The returned Column yields a struct ``{status, response, error_code,
+    error_message}`` for each row. ``response`` is itself a struct
+    ``{status_code, body, headers}``. Use ``[...]`` to pick a field, e.g.
+    ``named_credential_request_col(...)["response"]["status_code"]``. A transport
+    failure sets ``status`` to ``ERROR`` and populates ``error_message`` (a non-2xx
+    HTTP response is still ``SUCCESS`` with its code in ``response.status_code``),
+    so a single bad row does not abort the whole Spark job.
+
+    Args:
+        request: The callout template — its symbolic reference, method, and
+            headers are applied to every row.
+        body: Optional per-row ``Column`` holding the request body as a
+            string (or null for no body).
+
+    Returns:
+        A Spark ``Column`` of ``StructType`` with fields ``status``,
+        ``response``, ``error_code``, and ``error_message``.
+    """
+    named_credential = Client()._get_spark_named_credential()
+    return named_credential.request_col(request, body=body)
+
+
 class DataCloudObjectType(Enum):
     DLO = "dlo"
     DMO = "dmo"
@@ -266,6 +319,14 @@ class _BaseClient:
         spark_llm_gateway: Optional custom :class:`SparkLLMGateway`.
         spark_einstein_predictions: Optional custom
             :class:`SparkEinsteinPredictions`.
+        spark_named_credential: Optional custom :class:`SparkNamedCredential`.
+
+    Example:
+    >>> client = Client()
+    >>> file_path = client.find_file_path("data.csv")
+    >>> dlo = client.read_dlo("my_dlo")
+    >>> client.write_to_dmo("my_dmo", dlo)
+    >>> answer = client.llm_gateway_generate_text("Generate a greeting message")
     """
 
     # Each concrete subclass gets its own ``_instance`` slot: reads fall through
@@ -285,6 +346,7 @@ class _BaseClient:
     _file: DefaultFindFilePath
     _spark_llm_gateway: Optional[SparkLLMGateway]
     _spark_einstein_predictions: Optional[SparkEinsteinPredictions]
+    _spark_named_credential: Optional[SparkNamedCredential]
     _data_layer_history: dict[DataCloudObjectType, set[str]]
     _code_type: str
 
@@ -295,13 +357,15 @@ def __new__(
         spark_provider: Optional[BaseSparkSessionProvider] = None,
         spark_llm_gateway: Optional[SparkLLMGateway] = None,
         spark_einstein_predictions: Optional[SparkEinsteinPredictions] = None,
+        spark_named_credential: Optional[SparkNamedCredential] = None,
         code_type: str = "script",
     ) -> _ClientT:
 
         if cls._instance is None:
             instance = super().__new__(cls)
             instance._spark_llm_gateway = spark_llm_gateway
             instance._spark_einstein_predictions = spark_einstein_predictions
+            instance._spark_named_credential = spark_named_credential
             # Initialize Readers and Writers from config
             # and/or provided reader and writer
             if reader is None or writer is None:
@@ -474,6 +538,41 @@ def _get_spark_einstein_predictions(self) -> SparkEinsteinPredictions:
             self._spark_einstein_predictions = _build_spark_einstein_predictions()
         return self._spark_einstein_predictions
 
+    def named_credential_request(
+        self,
+        request: "HTTPRequest",
+        body: Optional[str] = None,
+    ) -> "HTTPResponse":
+        """Issue a one-shot Named Credential external callout. This is the
+        scalar counterpart to :func:`named_credential_request_col`: it runs
+        **once** on the driver — not per row. Use the column helper method
+        instead when you want to fan a callout out across every row of a
+        DataFrame.
+
+        Example:
+
+            >>> from datacustomcode.named_credential.types.http_request_builder \\
+            ...     import HTTPRequestBuilder
+            >>> request = (
+            ...     HTTPRequestBuilder().set_url("callout:NC/search").build()
+            ... )
+            >>> response = Client().named_credential_request(request)
+
+        Args:
+            request: The callout request
+            body: Optional request body. Set the ``Content-Type`` header to
+                match the format; the SDK does not assume or inject one.
+
+        Returns:
+            The external service's response.
+        """
+        return self._get_spark_named_credential().request(request, body=body)
+
+    def _get_spark_named_credential(self) -> SparkNamedCredential:
+        if self._spark_named_credential is None:
+            self._spark_named_credential = _build_spark_named_credential()
+        return self._spark_named_credential
+
     def _validate_data_layer_history_does_not_contain(
         self, data_cloud_object_type: DataCloudObjectType
     ) -> None:

src/datacustomcode/config.yaml

@@ -34,3 +34,9 @@ llm_gateway_config:
 
 spark_llm_gateway_config:
   type_config_name: DefaultSparkLLMGateway
+
+named_credential_config:
+  type_config_name: DefaultNamedCredential
+
+spark_named_credential_config:
+  type_config_name: DefaultSparkNamedCredential

src/datacustomcode/deploy.py

@@ -38,6 +38,9 @@
 
 from datacustomcode.cmd import cmd_output
 from datacustomcode.constants import REQUEST_TYPE_TO_FEATURE
+from datacustomcode.named_credential.direct.credentials import (
+    EXTERNAL_CALLOUT_CREDENTIAL,
+)
 from datacustomcode.scan import find_base_directory, get_package_type
 
 DATA_CUSTOM_CODE_PATH = "services/data/v63.0/ssot/data-custom-code"
@@ -248,6 +251,8 @@ def create_deployment(
 )
 PY_FILES_PATH = os.path.join("payload", "py-files")
 ZIP_FILE_NAME = "deployment.zip"
+# Local-only files that must never be packaged into the deployment zip.
+EXCLUDED_FILES = (".DS_Store", EXTERNAL_CALLOUT_CREDENTIAL)
 
 
 def prepare_dependency_archive(
@@ -628,9 +633,9 @@ def zip(
 
     with zipfile.ZipFile(ZIP_FILE_NAME, "w", zipfile.ZIP_DEFLATED) as zipf:
         for root, dirs, files in os.walk(directory):
-            # Skip .DS_Store files when adding to zip
+            # Skip .DS_Store and local credentials.
             for file in files:
-                if file != ".DS_Store":
+                if file not in EXCLUDED_FILES:
                     abs_path = os.path.join(root, file)
                     arcname = os.path.relpath(abs_path, directory)
                     zipf.write(abs_path, arcname)

src/datacustomcode/function/runtime.py

@@ -23,6 +23,8 @@
 from datacustomcode.function.base import BaseRuntime
 from datacustomcode.llm_gateway.base import LLMGateway
 from datacustomcode.llm_gateway_config import llm_gateway_config
+from datacustomcode.named_credential.base import NamedCredential
+from datacustomcode.named_credential_config import named_credential_config
 
 
 class Runtime(BaseRuntime):
@@ -69,6 +71,7 @@ def __init__(self) -> None:
         self._llm_gateway: Optional[LLMGateway] = None
         self._file = DefaultFindFilePath()
         self._einstein_predictions: Optional[EinsteinPredictions] = None
+        self._named_credential: Optional[NamedCredential] = None
 
     @property
     def llm_gateway(self) -> LLMGateway:
@@ -98,3 +101,16 @@ def einstein_predictions(self) -> EinsteinPredictions:
                 einstein_predictions_config.einstein_predictions_config.to_object()
             )
         return self._einstein_predictions
+
+    @property
+    def named_credential(self) -> NamedCredential:
+        if self._named_credential is None:
+            if named_credential_config.named_credential_config is None:
+                raise RuntimeError(
+                    "Named Credential is not configured. Add "
+                    "'named_credential_config' section to config.yaml"
+                )
+            self._named_credential = (
+                named_credential_config.named_credential_config.to_object()
+            )
+        return self._named_credential

src/datacustomcode/named_credential/__init__.py

@@ -0,0 +1,28 @@
+# 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.named_credential.base import NamedCredential
+from datacustomcode.named_credential.default import DefaultNamedCredential
+from datacustomcode.named_credential.errors import NamedCredentialCallError
+from datacustomcode.named_credential.spark_base import SparkNamedCredential
+from datacustomcode.named_credential.spark_default import DefaultSparkNamedCredential
+
+__all__ = [
+    "DefaultNamedCredential",
+    "DefaultSparkNamedCredential",
+    "NamedCredential",
+    "NamedCredentialCallError",
+    "SparkNamedCredential",
+]

src/datacustomcode/named_credential/base.py

@@ -0,0 +1,54 @@
+# 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 ABC, abstractmethod
+from typing import TYPE_CHECKING, Optional
+
+from datacustomcode.mixin import UserExtendableNamedConfigMixin
+
+if TYPE_CHECKING:
+    from datacustomcode.named_credential.types.http_request import HTTPRequest
+    from datacustomcode.named_credential.types.http_response import HTTPResponse
+
+
+class NamedCredential(ABC, UserExtendableNamedConfigMixin):
+    CONFIG_NAME: str
+
+    def __init__(self, **kwargs):
+        pass
+
+    @abstractmethod
+    def request(
+        self,
+        request: HTTPRequest,
+        body: Optional[str] = None,
+    ) -> HTTPResponse:
+        """Make an external callout through a Named Credential.
+
+        The endpoint and its authentication are resolved server-side from the
+        Named Credential referenced by ``request.url``; the function never sees
+        the external credential.
+
+        Args:
+            request: The callout request
+            body: Optional request body. Set the
+                ``Content-Type`` header to match the format; the SDK
+                does not assume or inject one.
+
+        Returns:
+            The external service's response.
+        """
+        ...