diff --git a/pyproject.toml b/pyproject.toml index 2d6938b..c0e868b 100644 --- a/pyproject.toml +++ b/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"] diff --git a/src/datacustomcode/client.py b/src/datacustomcode/client.py index b43e423..814ac5a 100644 --- a/src/datacustomcode/client.py +++ b/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): diff --git a/src/datacustomcode/named_credential/spark_base.py b/src/datacustomcode/named_credential/spark_base.py index 2c47743..8e46457 100644 --- a/src/datacustomcode/named_credential/spark_base.py +++ b/src/datacustomcode/named_credential/spark_base.py @@ -66,11 +66,12 @@ 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. @@ -78,6 +79,11 @@ def request_col( 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 diff --git a/src/datacustomcode/named_credential/spark_default.py b/src/datacustomcode/named_credential/spark_default.py index 8f3ea23..3e13fcc 100644 --- a/src/datacustomcode/named_credential/spark_default.py +++ b/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( diff --git a/tests/spark/test_column_hints_deferred.py b/tests/spark/test_column_hints_deferred.py index edf0622..9c094eb 100644 --- a/tests/spark/test_column_hints_deferred.py +++ b/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( diff --git a/tests/test_client.py b/tests/test_client.py index 4f3f2f4..fdcdaa8 100644 --- a/tests/test_client.py +++ b/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: diff --git a/tests/test_named_credential.py b/tests/test_named_credential.py index d486a74..70a7399 100644 --- a/tests/test_named_credential.py +++ b/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."""