Support per-row path in named_credential_request_col (#158)

* Support per-row url in named_credential_request_col

Add an optional url Column to named_credential_request_col /
SparkNamedCredential.request_col so each row can call a different endpoint
(e.g. per-entity REST GET lookups). A null row value falls back to
request.url; omitting url keeps the existing expression unchanged.

* Lint fixes

* Take a per-row path suffix instead of a full url

Rows may only select an endpoint under the template's Named Credential: the
path is appended to request.url, with '/' inserted unless it starts with '/'
or '?', so a row can never change the Named Credential name. A null or empty
path uses request.url.

Author: rkang-sf

pyproject.toml

@@ -87,10 +87,10 @@ warn_unused_configs = true
 
 [tool.poetry]
 include = [
-  {path = "src/datacustomcode/templates/**/*"},
-  {path = "src/datacustomcode/config.yaml"}
+  { path = "src/datacustomcode/templates/**/*" },
+  { path = "src/datacustomcode/config.yaml" }
 ]
-packages = [{include = "datacustomcode", from = "src"}]
+packages = [{ include = "datacustomcode", from = "src" }]
 version = "0.0.0"
 
 [tool.poetry.build]
@@ -115,7 +115,7 @@ coverage = ">=7.0.0,<8.0.0"
 ipykernel = "^6.29.5"
 mypy = "*"
 pipreqs = "*"
-poetry-dynamic-versioning = {extras = ["plugin"], version = "^1.8.2"}
+poetry-dynamic-versioning = { extras = ["plugin"], version = "^1.8.2" }
 pre-commit = "*"
 pytest = "*"
 pytest-cov = "*"
@@ -127,7 +127,7 @@ types-requests = "*"
 [tool.poetry.plugins."poetry.plugin"]
 
 [tool.poetry.requires-plugins]
-poetry-dynamic-versioning = {version = ">=1.0.0,<2.0.0", extras = ["plugin"]}
+poetry-dynamic-versioning = { version = ">=1.0.0,<2.0.0", extras = ["plugin"] }
 
 [tool.poetry.scripts]
 datacustomcode = "datacustomcode.cli:cli"
@@ -152,7 +152,6 @@ testpaths = ["tests"]
 [tool.ruff]
 fix = true
 line-length = 88
-target-version = 'py310'
 lint.ignore = [
   # do not assign a lambda expression, use a def
   'E731',
@@ -211,6 +210,7 @@ lint.select = [
   'RUF',
   'S102'
 ]
+target-version = 'py310'
 
 [tool.setuptools.package_data]
 datacustomcode = ["templates/**/*", "config.yaml"]

src/datacustomcode/client.py

@@ -231,11 +231,12 @@ def einstein_predict_col(
 def named_credential_request_col(
     request: "HTTPRequest",
     body: Optional["Column"] = None,
+    path: 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
+    The method and headers are fixed for the call (taken from ``request``);
+    ``body`` and, optionally, ``path`` vary 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.
@@ -253,13 +254,20 @@ def named_credential_request_col(
             headers are applied to every row.
         body: Optional per-row ``Column`` holding the request body as a
             string (or null for no body).
+        path: Optional per-row ``Column`` appended to ``request.url``, e.g.
+            ``concat(lit("?address="), col("address"))`` or ``lit("/v1/models")``.
+            The callout always stays on ``request.url``'s Named Credential: a
+            path not starting with ``/`` or ``?`` is joined with ``/``. A null
+            or empty row value uses ``request.url`` as is.
 
     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)
+    if path is None:
+        return named_credential.request_col(request, body=body)
+    return named_credential.request_col(request, body=body, path=path)
 
 
 class DataCloudObjectType(Enum):

src/datacustomcode/named_credential/spark_base.py

@@ -66,18 +66,24 @@ def request_col(
         self,
         request: HTTPRequest,
         body: Optional["Column"] = None,
+        path: 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
+        The method and headers are fixed for the call (taken from ``request``);
+        ``body`` and, optionally, ``path`` vary 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).
+            path: Optional per-row ``Column`` appended to ``request.url``
+                (e.g. ``/v1/models`` or ``?id=42``) so each row can target a
+                different endpoint of the same Named Credential. A path not
+                starting with ``/`` or ``?`` is joined with ``/``; a null or
+                empty row value uses ``request.url`` as is.
 
         Returns:
             A ``Column`` yielding a struct

src/datacustomcode/named_credential/spark_default.py

@@ -79,6 +79,7 @@ def request_col(
         self,
         request: "HTTPRequest",
         body: Optional["Column"] = None,
+        path: Optional["Column"] = None,
     ) -> "Column":
         """Per-row callout via a client-side Spark UDF.
 
@@ -114,11 +115,46 @@ def request_col(
             ]
         )
 
-        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)
+
+        if path is None:
+
+            def _callout(body_str: Optional[str]) -> Dict[str, Any]:
+                return _invoke_callout_as_struct(
+                    self._named_credential, request, body_str
+                )
+
+            return udf(_callout, result_schema)(body_col)
+
+        def _callout_with_path(
+            body_str: Optional[str], path_str: Optional[str]
+        ) -> Dict[str, Any]:
+            row_request = (
+                request.model_copy(
+                    update={"url": join_callout_path(request.url, path_str)}
+                )
+                if path_str
+                else request
+            )
+            return _invoke_callout_as_struct(
+                self._named_credential, row_request, body_str
+            )
+
+        return udf(_callout_with_path, result_schema)(body_col, path)
+
+
+def join_callout_path(base_url: str, path: Optional[str]) -> str:
+    """Append a per-row ``path`` to the template ``base_url``.
+
+    A path not starting with ``/`` or ``?`` is joined with ``/`` so the row can
+    never extend the Named Credential name (``callout:NC`` + ``2/x`` must not
+    become ``callout:NC2/x``). A null or empty path returns ``base_url``.
+    """
+    if not path:
+        return base_url
+    if path.startswith(("/", "?")):
+        return base_url + path
+    return f"{base_url}/{path}"
 
 
 def _invoke_callout_as_struct(

tests/spark/test_column_hints_deferred.py

@@ -38,7 +38,8 @@ def _run(script: str) -> None:
         [sys.executable, "-c", textwrap.dedent(script)],
         capture_output=True,
         text=True,
-        timeout=60, check=False,
+        timeout=60,
+        check=False,
     )
     if result.returncode != 0:
         pytest.fail(

tests/test_client.py

@@ -733,6 +733,27 @@ def test_delegates_to_spark_named_credential(self, mock_build, reset_client):
         assert result is sentinel_col
         mock_nc.request_col.assert_called_once_with(request, body=body_col)
 
+    @patch("datacustomcode.client._build_spark_named_credential")
+    def test_forwards_per_row_path_column(self, mock_build, reset_client):
+        mock_nc = MagicMock()
+        sentinel_col = MagicMock(name="col")
+        mock_nc.request_col.return_value = sentinel_col
+        mock_build.return_value = mock_nc
+
+        reader = MagicMock(spec=BaseDataCloudReader)
+        writer = MagicMock(spec=BaseDataCloudWriter)
+        Client(reader=reader, writer=writer)
+
+        request = MagicMock(name="request")
+        body_col = MagicMock(name="body_col")
+        path_col = MagicMock(name="path_col")
+        result = named_credential_request_col(request, body_col, path=path_col)
+
+        assert result is sentinel_col
+        mock_nc.request_col.assert_called_once_with(
+            request, body=body_col, path=path_col
+        )
+
 
 class TestClientNamedCredentialRequest:
 

tests/test_named_credential.py

@@ -403,6 +403,138 @@ def test_defaults_body_to_typed_null_column(self, mock_lit, mock_udf):
         # With no body column, a typed null string column is applied instead.
         sentinel_udf.assert_called_once_with(null_col)
 
+    @patch("pyspark.sql.functions.udf")
+    @patch("pyspark.sql.functions.lit")
+    def test_per_row_path_is_appended_to_request_url(self, mock_lit, mock_udf):
+        from datacustomcode.named_credential.spark_default import (
+            DefaultSparkNamedCredential,
+        )
+
+        sentinel_udf = MagicMock(name="udf")
+        mock_udf.return_value = sentinel_udf
+
+        underlying = MagicMock()
+        underlying.request.return_value = HTTPResponse(
+            status_code=200, headers={}, body="{}"
+        )
+        spark_nc = DefaultSparkNamedCredential(named_credential=underlying)
+
+        request = (
+            HTTPRequestBuilder()
+            .set_url("callout:gemini")
+            .set_method("POST")
+            .set_headers({"Accept": "application/json"})
+            .build()
+        )
+        body_col = MagicMock(name="body_col")
+        path_col = MagicMock(name="path_col")
+        spark_nc.request_col(request, body_col, path=path_col)
+
+        # The UDF is applied over both the body and the per-row path columns.
+        sentinel_udf.assert_called_once_with(body_col, path_col)
+
+        udf_fn = mock_udf.call_args.args[0]
+        out = udf_fn('{"q": 1}', "/gemini-2.5-pro")
+
+        assert out["status"] == "SUCCESS"
+        sent_request, sent_body = underlying.request.call_args.args
+        # The row's path is joined onto the template url; method/headers kept.
+        assert sent_request.url == "callout:gemini/gemini-2.5-pro"
+        assert sent_request.method == "POST"
+        assert sent_request.headers == {"Accept": "application/json"}
+        assert sent_body == '{"q": 1}'
+        # The caller's template request is left untouched.
+        assert request.url == "callout:gemini"
+
+    @pytest.mark.parametrize(
+        "base, path, expected",
+        [
+            ("callout:gemini", "/v1/models", "callout:gemini/v1/models"),
+            ("callout:gemini", "?key=a", "callout:gemini?key=a"),
+            (
+                "callout:Geo/geocode",
+                "?address=1+Main",
+                "callout:Geo/geocode?address=1+Main",
+            ),
+            ("callout:Geo/v1", "/places/42", "callout:Geo/v1/places/42"),
+            # Without a leading separator a "/" is inserted, so the row can never
+            # extend the Named Credential name (callout:gemini2/...).
+            ("callout:gemini", "2/v1", "callout:gemini/2/v1"),
+            ("callout:gemini", "@evil.example.com", "callout:gemini/@evil.example.com"),
+        ],
+    )
+    @patch("pyspark.sql.functions.udf")
+    @patch("pyspark.sql.functions.lit")
+    def test_path_join_keeps_named_credential_fixed(
+        self, mock_lit, mock_udf, base, path, expected
+    ):
+        from datacustomcode.named_credential.spark_default import (
+            DefaultSparkNamedCredential,
+        )
+
+        mock_udf.return_value = MagicMock(name="udf")
+        underlying = MagicMock()
+        underlying.request.return_value = HTTPResponse(
+            status_code=200, headers={}, body=""
+        )
+        spark_nc = DefaultSparkNamedCredential(named_credential=underlying)
+
+        request = HTTPRequestBuilder().set_url(base).build()
+        spark_nc.request_col(request, path=MagicMock(name="path_col"))
+
+        udf_fn = mock_udf.call_args.args[0]
+        udf_fn(None, path)
+
+        assert underlying.request.call_args.args[0].url == expected
+
+    @pytest.mark.parametrize("path", [None, ""])
+    @patch("pyspark.sql.functions.udf")
+    @patch("pyspark.sql.functions.lit")
+    def test_null_or_empty_row_path_uses_request_url(self, mock_lit, mock_udf, path):
+        from datacustomcode.named_credential.spark_default import (
+            DefaultSparkNamedCredential,
+        )
+
+        mock_udf.return_value = MagicMock(name="udf")
+
+        underlying = MagicMock()
+        underlying.request.return_value = HTTPResponse(
+            status_code=200, headers={}, body=""
+        )
+        spark_nc = DefaultSparkNamedCredential(named_credential=underlying)
+
+        request = HTTPRequestBuilder().set_url("callout:NC/default").build()
+        spark_nc.request_col(request, path=MagicMock(name="path_col"))
+
+        udf_fn = mock_udf.call_args.args[0]
+        udf_fn('{"a": 1}', path)
+
+        sent_request, sent_body = underlying.request.call_args.args
+        assert sent_request is request
+        assert sent_body == '{"a": 1}'
+
+    @patch("pyspark.sql.functions.udf")
+    @patch("pyspark.sql.functions.lit")
+    def test_path_without_body_applies_typed_null_body(self, mock_lit, mock_udf):
+        from datacustomcode.named_credential.spark_default import (
+            DefaultSparkNamedCredential,
+        )
+
+        null_col = MagicMock(name="null_col")
+        lit_none = MagicMock(name="lit_none")
+        lit_none.cast.return_value = null_col
+        mock_lit.return_value = lit_none
+
+        sentinel_udf = MagicMock(name="udf")
+        mock_udf.return_value = sentinel_udf
+
+        spark_nc = DefaultSparkNamedCredential(named_credential=MagicMock())
+        request = HTTPRequestBuilder().set_url("callout:NC/status").build()
+        path_col = MagicMock(name="path_col")
+        spark_nc.request_col(request, path=path_col)
+
+        sentinel_udf.assert_called_once_with(null_col, path_col)
+
 
 class TestInvokeCalloutAsStruct:
     """The callout-to-struct shaping shared by every row."""