salesforce-data-customcode 6.1.0.dev1__py3-none-any.whl → 6.1.0.dev2__py3-none-any.whl

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 (
+    CredentialError,
+    CredentialStore,
+)
+from datacustomcode.named_credential.direct.url_resolver import resolve_base_url
+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 CredentialError(
+                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 CredentialError(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/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"

datacustomcode/named_credential/types/http_request.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 enum import Enum
+from typing import Dict
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    field_validator,
+)
+
+from datacustomcode.named_credential.types.http_method import HTTPMethod
+
+
+class HTTPRequest(BaseModel):
+    """External callout request. The endpoint and its auth are resolved
+    server-side from the Named Credential referenced by ``url``, which uses
+    ``callout:<NamedCredential>/<path>`` syntax."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    url: str = Field(
+        ...,
+        min_length=1,
+        description="Symbolic Named Credential reference, "
+        "e.g. 'callout:<NamedCredential>/<path>'",
+    )
+    method: str = Field(
+        default="GET",
+        description="HTTP method (GET, POST, PUT, DELETE, PATCH, HEAD, OPTIONS)",
+    )
+    headers: Dict[str, str] = Field(default_factory=dict, description="Request headers")
+
+    @field_validator("method", mode="before")
+    @classmethod
+    def _normalize_method(cls, value: object) -> str:
+        # Accept str, this module's HTTPMethod, or http.HTTPMethod (3.11+).
+        if isinstance(value, Enum):
+            method = str(value.value)
+        else:
+            method = str(value)
+        method = method.upper()
+        if method not in {m.value for m in HTTPMethod}:
+            supported = ", ".join(m.value for m in HTTPMethod)
+            raise ValueError(
+                f"Unsupported HTTP method '{method}'. "
+                f"Named Credential callouts support {supported}."
+            )
+        return method

datacustomcode/named_credential/types/http_request_builder.py

@@ -0,0 +1,55 @@
+# 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 typing import Dict, Union
+
+from datacustomcode.named_credential.types.http_method import HTTPMethod
+from datacustomcode.named_credential.types.http_request import HTTPRequest
+
+
+class HTTPRequestBuilder:
+    def __init__(self) -> None:
+        self._url = ""
+        self._method: Union[str, HTTPMethod] = HTTPMethod.GET
+        self._headers: Dict[str, str] = {}
+
+    def set_url(self, url: str) -> "HTTPRequestBuilder":
+        """Set the symbolic Named Credential reference.
+
+        Args:
+            url: e.g. 'callout:<NamedCredential>/<path>'
+        """
+        self._url = url
+        return self
+
+    def set_method(self, method: Union[str, HTTPMethod]) -> "HTTPRequestBuilder":
+        """Set the HTTP method.
+
+        Accepts this module's ``HTTPMethod``, ``http.HTTPMethod`` (Python 3.11+),
+        or a plain string such as ``"GET"``.
+        """
+        self._method = method
+        return self
+
+    def set_headers(self, headers: Dict[str, str]) -> "HTTPRequestBuilder":
+        self._headers = headers
+        return self
+
+    def build(self) -> HTTPRequest:
+        return HTTPRequest(
+            url=self._url,
+            method=self._method,  # type: ignore[arg-type]
+            headers=self._headers,
+        )

datacustomcode/named_credential/types/http_response.py

@@ -0,0 +1,43 @@
+# 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 typing import Dict
+
+from pydantic import BaseModel, Field
+
+
+class HTTPResponse(BaseModel):
+    """Response from a Named Credential external callout."""
+
+    status_code: int = Field(..., description="HTTP status code", ge=0)
+    headers: Dict[str, str] = Field(
+        default_factory=dict, description="Response headers"
+    )
+    body: str = Field(
+        default="",
+        description="Raw response body, verbatim (any format). The SDK does not "
+        "parse it; the caller decodes as needed (e.g. json.loads). Empty string "
+        "if the response had no body.",
+    )
+
+    @property
+    def is_success(self) -> bool:
+        """Check if the request succeeded (2xx)."""
+        return 200 <= self.status_code < 300
+
+    @property
+    def is_error(self) -> bool:
+        """Check if the request failed."""
+        return not self.is_success