salesforce-data-customcode 6.1.0.dev5__py3-none-any.whl → 6.1.0.dev7__py3-none-any.whl

datacustomcode/named_credential/direct/credentials.py

@@ -0,0 +1,122 @@
+# 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.
+"""Load developer credentials for direct callouts from ``external_callout_config.json``.
+
+The file lives in the parent of the payload folder so it is never packaged in
+the deployment zip. Its ``credentials`` section is keyed by callout reference
+(``callout:<NamedCredential>``); each value carries a mandatory ``auth_type``
+and an optional ``target_url``.
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+import json
+import os
+from pathlib import Path
+from typing import (
+    Any,
+    Dict,
+    Optional,
+)
+
+# Default file name; discovered in the parent of the payload folder.
+EXTERNAL_CALLOUT_CREDENTIAL = "external_callout_config.json"
+# Absolute-path override, primarily for tests and non-standard layouts.
+CREDENTIAL_FILE_ENV_VAR = "DATACUSTOMCODE_EXTERNAL_CALLOUT_CONFIG"
+
+
+class AuthType(str, Enum):
+    """External Credential auth types supported by External Services."""
+
+    AWS_SIG_V4 = "AwsSv4"
+    BASIC = "Basic"
+    CUSTOM = "Custom"
+    JWT = "Jwt"
+    OAUTH = "OAuth"
+
+
+class CredentialError(RuntimeError):
+    """Raised when credentials cannot be found or are invalid."""
+
+
+def _discover_credential_file() -> Optional[Path]:
+    """Find the config file via env override, then by walking up from cwd."""
+    override = os.environ.get(CREDENTIAL_FILE_ENV_VAR)
+    if override:
+        return Path(override)
+
+    for directory in (Path.cwd(), *Path.cwd().parents):
+        candidate = directory / EXTERNAL_CALLOUT_CREDENTIAL
+        if candidate.is_file():
+            return candidate
+    return None
+
+
+class CredentialStore:
+    """Reads ``external_callout_config.json`` and returns per-callout config."""
+
+    def __init__(self, credential_file: Optional[str] = None) -> None:
+        self._explicit_path = Path(credential_file) if credential_file else None
+        self._credentials: Optional[Dict[str, Dict[str, Any]]] = None
+
+    def _load(self) -> Dict[str, Dict[str, Any]]:
+        if self._credentials is not None:
+            return self._credentials
+
+        path = self._explicit_path or _discover_credential_file()
+        if path is None or not path.is_file():
+            raise CredentialError(
+                f"Could not find '{EXTERNAL_CALLOUT_CREDENTIAL}'. Place it in the "
+                f"parent of your payload folder, or set "
+                f"${CREDENTIAL_FILE_ENV_VAR} to its path."
+            )
+        try:
+            with open(path, "r") as f:
+                data = json.load(f)
+        except (OSError, json.JSONDecodeError) as exc:
+            raise CredentialError(f"Failed to read '{path}': {exc}") from exc
+
+        # Per-callout entries live under the ``credentials`` section, leaving
+        # room for other config sections alongside them in the future.
+        credentials = data.get("credentials") if isinstance(data, dict) else None
+        if not isinstance(credentials, dict):
+            raise CredentialError(
+                f"'{path}' must be a JSON object with a 'credentials' section "
+                f"keyed by callout reference."
+            )
+        self._credentials = credentials
+        return credentials
+
+    def get(self, callout_key: str) -> Dict[str, Any]:
+        """Return the config for a callout key (e.g. ``callout:AWS_S3_Service``).
+
+        Raises:
+            CredentialError: if the key is missing or has no ``auth_type``.
+        """
+        credentials = self._load()
+        config = credentials.get(callout_key)
+        if config is None:
+            raise CredentialError(
+                f"No credential configuration found for '{callout_key}'. "
+                f"Add it to the 'credentials' section of "
+                f"'{EXTERNAL_CALLOUT_CREDENTIAL}'."
+            )
+        if not isinstance(config, dict) or not config.get("auth_type"):
+            raise CredentialError(
+                f"Credential for '{callout_key}' is missing the mandatory "
+                f"'auth_type' field."
+            )
+        return config

datacustomcode/named_credential/direct/transport.py

@@ -0,0 +1,110 @@
+# 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.
+"""Send a Named Credential callout over HTTP.
+
+Resolves the ``callout:<NamedCredential>/<path>`` reference to a real endpoint
+(via the Named Credential Connect API, falling back to ``target_url`` in
+``external_callout_config.json``), attaches the credential's auth, and returns the raw
+``{status_code, headers, body}`` response.
+"""
+
+from __future__ import annotations
+
+from typing import (
+    Any,
+    Dict,
+    Optional,
+)
+
+import requests
+
+from datacustomcode.named_credential.direct.auth import DynamicAuthHandler
+from datacustomcode.named_credential.direct.credentials import CredentialStore
+from datacustomcode.named_credential.direct.url_resolver import resolve_base_url
+from datacustomcode.named_credential.errors import NamedCredentialCallError
+from datacustomcode.token_provider import (
+    CredentialsTokenProvider,
+    SFCLITokenProvider,
+    TokenProvider,
+)
+
+
+class DirectCalloutTransport:
+    def __init__(
+        self,
+        credentials_profile: str = "default",
+        sf_cli_org: Optional[str] = None,
+        credential_file: Optional[str] = None,
+    ) -> None:
+        self._store = CredentialStore(credential_file)
+        self._token_provider = self._build_token_provider(
+            credentials_profile, sf_cli_org
+        )
+        # Resolved base URL per callout key. Stable for the transport's life, so
+        # cache it to avoid a token fetch + Connect API call on every row
+        self._base_url_cache: Dict[str, str] = {}
+
+    @staticmethod
+    def _build_token_provider(
+        credentials_profile: str, sf_cli_org: Optional[str]
+    ) -> TokenProvider:
+        if sf_cli_org:
+            return SFCLITokenProvider(sf_cli_org)
+        return CredentialsTokenProvider(credentials_profile)
+
+    def callout(self, callout_request: Dict[str, Any]) -> Dict[str, Any]:
+        raw_url = callout_request["path"]
+        if not raw_url.startswith("callout:"):
+            raise NamedCredentialCallError(
+                f"Callout URL must start with 'callout:', got '{raw_url}'."
+            )
+
+        # Split the named credential reference at the first '/' or '?'; the remainder
+        # path or query string is appended to the resolved base URL verbatim.
+        sep_idx = min(
+            (i for i in (raw_url.find("/"), raw_url.find("?")) if i != -1),
+            default=len(raw_url),
+        )
+        callout_key = raw_url[:sep_idx]
+        path_suffix = raw_url[sep_idx:]
+
+        if callout_key == "callout:":
+            raise NamedCredentialCallError(
+                f"Named Credential name is empty in URL '{raw_url}'."
+            )
+
+        cred_config = self._store.get(callout_key)
+        base_url = self._base_url_cache.get(callout_key)
+        if base_url is None:
+            base_url = resolve_base_url(callout_key, cred_config, self._token_provider)
+            self._base_url_cache[callout_key] = base_url
+
+        body = callout_request.get("body") or None
+        # Headers are passed; the SDK assumes no Content-Type.
+        headers = dict(callout_request.get("headers", {}))
+
+        response = requests.request(
+            method=callout_request["method"],
+            url=base_url + path_suffix,
+            headers=headers,
+            data=body,
+            auth=DynamicAuthHandler(cred_config),
+            timeout=30,
+        )
+        return {
+            "status_code": response.status_code,
+            "headers": dict(response.headers),
+            "body": response.text,
+        }

datacustomcode/named_credential/direct/url_resolver.py

@@ -0,0 +1,112 @@
+# 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.
+"""Resolve a ``callout:<NamedCredential>`` reference to a real base URL.
+
+Prefers the NamedCredential Connect API ``calloutUrl``; falls back to the
+developer-supplied ``target_url`` when the API is unavailable or the Named
+Credential is not yet onboarded.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Optional,
+)
+
+import requests
+
+from datacustomcode.named_credential.direct.credentials import CredentialError
+
+if TYPE_CHECKING:
+    from datacustomcode.token_provider import TokenProvider
+
+logger = logging.getLogger(__name__)
+
+NAMED_CREDENTIAL_PATH = (
+    "services/data/v63.0/named-credentials/named-credential-setup/{name}"
+)
+
+
+def _callout_url_from_connect_api(
+    developer_name: str, token_provider: TokenProvider
+) -> Optional[str]:
+    """Fetch ``calloutUrl`` for a Named Credential, or None on any failure."""
+    try:
+        token = token_provider.get_token()
+    except Exception as exc:
+        # No usable token (e.g. not logged in) — expected; fall back quietly.
+        logger.debug("Could not obtain a token for %s: %s", developer_name, exc)
+        return None
+
+    try:
+        path = NAMED_CREDENTIAL_PATH.format(name=developer_name)
+        url = f"{token.instance_url.rstrip('/')}/{path}"
+        response = requests.get(
+            url,
+            headers={"Authorization": f"Bearer {token.access_token}"},
+            timeout=30,
+        )
+        response.raise_for_status()
+        callout_url = response.json().get("calloutUrl")
+        return callout_url or None
+    except Exception as exc:
+        # A token was obtained but the Connect API call failed — likely a
+        # misconfiguration (named credential not onboarded, missing
+        # permission, wrong org). Surface it before falling back to target_url.
+        logger.warning(
+            "Connect API URL resolution failed for %s: %s. "
+            "Falling back to 'target_url' from external_callout_config.json if set.",
+            developer_name,
+            exc,
+        )
+        return None
+
+
+def resolve_base_url(
+    callout_key: str,
+    cred_config: Dict[str, Any],
+    token_provider: Optional[TokenProvider],
+) -> str:
+    """Resolve the base URL for a callout key.
+
+    Args:
+        callout_key: e.g. ``callout:Nominatim_Geocoding``.
+        cred_config: The callout's ``external_callout_config.json`` entry.
+        token_provider: Provides a token/instance URL for the Connect API; when
+            None, only ``target_url`` is used.
+
+    Raises:
+        CredentialError: if no URL can be resolved.
+    """
+    developer_name = callout_key.split(":", 1)[1] if ":" in callout_key else callout_key
+
+    base_url: Optional[str] = None
+    if token_provider is not None:
+        base_url = _callout_url_from_connect_api(developer_name, token_provider)
+
+    if not base_url:
+        base_url = cred_config.get("target_url") or None
+
+    if not base_url:
+        raise CredentialError(
+            f"Could not resolve a URL for '{callout_key}'. Ensure the Named "
+            f"Credential exists (sf login) or set 'target_url' in "
+            f"external_callout_config.json."
+        )
+    return base_url.rstrip("/")

datacustomcode/named_credential/errors.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.
+"""Exceptions raised by Named Credential implementations."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+
+class NamedCredentialCallError(RuntimeError):
+    """Raised when a Named Credential external callout fails."""
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        status: Optional[object] = None,
+        error_code: Optional[str] = None,
+        error_message: Optional[str] = None,
+    ) -> None:
+        super().__init__(message)
+        self.status = status
+        self.error_code = error_code
+        self.error_message = error_message

datacustomcode/named_credential/spark_base.py

@@ -0,0 +1,93 @@
+# 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,
+    Any,
+    Optional,
+)
+
+from datacustomcode.mixin import UserExtendableNamedConfigMixin
+
+if TYPE_CHECKING:
+    from pyspark.sql import Column
+
+    from datacustomcode.named_credential.types.http_request import HTTPRequest
+    from datacustomcode.named_credential.types.http_response import HTTPResponse
+
+
+class SparkNamedCredential(ABC, UserExtendableNamedConfigMixin):
+    """Named Credential external callout for script (Spark) code.
+
+    The callout is a one-shot request that runs on the driver. The endpoint and
+    its authentication are resolved from the Named Credential referenced by
+    ``request.url``.
+    """
+
+    CONFIG_NAME: str
+
+    def __init__(self, **kwargs: Any) -> None:
+        pass
+
+    @abstractmethod
+    def request(
+        self,
+        request: HTTPRequest,
+        body: Optional[str] = None,
+    ) -> HTTPResponse:
+        """Make an external callout through a Named Credential.
+
+        Args:
+            request: The callout request
+            body: Optional request body. Set the
+                ``Content-Type`` header yourself; the SDK does not assume one.
+
+        Returns:
+            The external service's response.
+        """
+        ...
+
+    @abstractmethod
+    def request_col(
+        self,
+        request: HTTPRequest,
+        body: Optional["Column"] = None,
+    ) -> "Column":
+        """Build a Spark ``Column`` that makes one external 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:`request` when the callout runs across a DataFrame so each row is
+        dispatched independently rather than one-shot on the driver.
+
+        Args:
+            request: The callout template
+            body: Optional per-row ``Column`` holding the request body as a
+                string, sent verbatim (or null for no body).
+
+        Returns:
+            A ``Column`` yielding a struct
+            ``{status, response, error_code, error_message}``. ``response`` is
+            itself a struct ``{status_code, body, headers}`` carrying the callout
+            response. Select a field with ``[...]``, e.g.
+            ``request_col(...)["response"]["status_code"]``. Returning a struct
+            means a single failing row does not abort the Spark job — a transport
+            failure sets ``status`` to ``ERROR`` with ``error_message``, while a
+            non-2xx HTTP response stays ``SUCCESS`` with its code in
+            ``response.status_code``.
+        """
+        ...

datacustomcode/named_credential/spark_default.py

@@ -0,0 +1,154 @@
+# 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 (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Optional,
+)
+
+from datacustomcode.named_credential.spark_base import SparkNamedCredential
+
+if TYPE_CHECKING:
+    from pyspark.sql import Column
+
+    from datacustomcode.named_credential.base import NamedCredential
+    from datacustomcode.named_credential.types.http_request import HTTPRequest
+    from datacustomcode.named_credential.types.http_response import HTTPResponse
+
+
+_STATUS_SUCCESS = "SUCCESS"
+_STATUS_ERROR = "ERROR"
+
+
+def _build_underlying_named_credential() -> "NamedCredential":
+    """Build the callout object from the configured ``named_credential_config``.
+
+    Raises ``RuntimeError`` if no ``named_credential_config`` section is set.
+    """
+    from datacustomcode.named_credential_config import named_credential_config
+
+    cfg = named_credential_config.named_credential_config
+    if cfg is None:
+        raise RuntimeError(
+            "named_credential_config is not configured. Add a "
+            "'named_credential_config' section to config.yaml."
+        )
+    return cfg.to_object()
+
+
+class DefaultSparkNamedCredential(SparkNamedCredential):
+    """
+    Callout for Spark, delegating to the shared implementation.
+    """
+
+    CONFIG_NAME = "DefaultSparkNamedCredential"
+
+    def __init__(
+        self,
+        named_credential: Optional["NamedCredential"] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        if named_credential is None:
+            named_credential = _build_underlying_named_credential()
+        self._named_credential: "NamedCredential" = named_credential
+
+    def request(
+        self,
+        request: "HTTPRequest",
+        body: Optional[str] = None,
+    ) -> "HTTPResponse":
+        return self._named_credential.request(request, body)
+
+    def request_col(
+        self,
+        request: "HTTPRequest",
+        body: Optional["Column"] = None,
+    ) -> "Column":
+        """Per-row callout via a client-side Spark UDF.
+
+        Returns a struct ``{status, response, error_code, error_message}`` where
+        ``response`` is itself a struct
+        ``{status_code, body, headers}`` carrying the callout's HTTP response, so
+        a script that selects ``request_col(...)["response"]["status_code"]``
+        behaves the same during development and in the Data Cloud runtime. Per-row
+        failures populate the error fields instead of aborting the Spark job.
+        """
+        from pyspark.sql.functions import lit, udf
+        from pyspark.sql.types import (
+            IntegerType,
+            MapType,
+            StringType,
+            StructField,
+            StructType,
+        )
+
+        http_response_schema = StructType(
+            [
+                StructField("status_code", IntegerType(), True),
+                StructField("body", StringType(), True),
+                StructField("headers", MapType(StringType(), StringType()), True),
+            ]
+        )
+        result_schema = StructType(
+            [
+                StructField("status", StringType(), True),
+                StructField("response", http_response_schema, True),
+                StructField("error_code", StringType(), True),
+                StructField("error_message", StringType(), True),
+            ]
+        )
+
+        def _callout(body_str: Optional[str]) -> Dict[str, Any]:
+            return _invoke_callout_as_struct(self._named_credential, request, body_str)
+
+        body_col = body if body is not None else lit(None).cast(StringType())
+        return udf(_callout, result_schema)(body_col)
+
+
+def _invoke_callout_as_struct(
+    named_credential: "NamedCredential",
+    request: "HTTPRequest",
+    body_str: Optional[str],
+) -> Dict[str, Any]:
+    """Run one callout and shape it into the shared result struct.
+
+    ``response`` is a nested struct ``{status_code, body, headers}`` where
+    ``body`` is the external response verbatim. Transport errors become per-row
+    ERROR structs rather than aborting the job.
+    """
+    try:
+        response = named_credential.request(request, body_str)
+    except Exception as exc:  # surface any transport error per row
+        return {
+            "status": _STATUS_ERROR,
+            "response": None,
+            "error_code": None,
+            "error_message": str(exc),
+        }
+
+    return {
+        "status": _STATUS_SUCCESS,
+        "response": {
+            "status_code": response.status_code,
+            "body": response.body,
+            "headers": response.headers,
+        },
+        "error_code": None,
+        "error_message": None,
+    }

datacustomcode/named_credential/types/__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/named_credential/types/http_method.py

@@ -0,0 +1,29 @@
+# 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 enum import Enum
+
+
+class HTTPMethod(str, Enum):
+    """Callout methods (GET, POST, PUT, DELETE, PATCH, HEAD, OPTIONS); usable on
+    Python 3.10 where ``http.HTTPMethod`` is unavailable."""
+
+    GET = "GET"
+    POST = "POST"
+    PUT = "PUT"
+    DELETE = "DELETE"
+    PATCH = "PATCH"
+    HEAD = "HEAD"
+    OPTIONS = "OPTIONS"