From bb9de27070fbdb84e988e9ffacd4a20b4b6fdbb8 Mon Sep 17 00:00:00 2001
From: Deepti Rajagopal <deepti.rajagopal@salesforce.com>
Date: Mon, 28 Sep 2026 08:12:42 -0700
Subject: [PATCH] Add QueryServiceUdf Function-shaped invoke option contract

Adds QueryServiceUdfV1Request/Response Pydantic models (following the
SearchIndexChunkingV1Request/Response pattern) plus the
REQUEST_TYPE_TO_FEATURE and USE_IN_FEATURE_MAPPING_FOR_CONNECT_API
entries needed for cli.py's infer_use_in_feature() to detect and tag
this feature at deploy time. No FEATURE_TEMPLATE_MAPPING scaffold yet
- see Known Gaps in ADD-INVOKE-OPTION-QueryServiceUdf.md.
---
 src/datacustomcode/constants.py               |  8 ++
 .../feature_types/query_service_udf.py        | 90 +++++++++++++++++++
 2 files changed, 98 insertions(+)
 create mode 100644 src/datacustomcode/function/feature_types/query_service_udf.py

diff --git a/src/datacustomcode/constants.py b/src/datacustomcode/constants.py
index 76b6a7c..295133a 100644
--- a/src/datacustomcode/constants.py
+++ b/src/datacustomcode/constants.py
@@ -36,10 +36,18 @@
 # Feature name to Connect API name mapping
 USE_IN_FEATURE_MAPPING_FOR_CONNECT_API = {
     "SearchIndexChunking": "UnstructuredChunking",
+    "QueryServiceUdf": "QueryServiceUdf",
 }
 
 # Pydantic request/response type names to feature names
 REQUEST_TYPE_TO_FEATURE = {
     "SearchIndexChunkingV1Request": "SearchIndexChunking",
     "SearchIndexChunkingV1Response": "SearchIndexChunking",
+    "QueryServiceUdfV1Request": "QueryServiceUdf",
+    "QueryServiceUdfV1Response": "QueryServiceUdf",
 }
+
+# NOTE: no FEATURE_TEMPLATE_MAPPING entry for QueryServiceUdf yet — that
+# requires a full `templates/function/query_service_udf/` scaffold
+# (Dockerfile, entrypoint, requirements) which is out of scope for this
+# enum-plumbing pass. See ADD-INVOKE-OPTION-QueryServiceUdf.md Known Gaps.
diff --git a/src/datacustomcode/function/feature_types/query_service_udf.py b/src/datacustomcode/function/feature_types/query_service_udf.py
new file mode 100644
index 0000000..7df7746
--- /dev/null
+++ b/src/datacustomcode/function/feature_types/query_service_udf.py
@@ -0,0 +1,90 @@
+# 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.
+
+"""
+Pydantic models for Query Service UDF V1
+
+PROVISIONAL: authored while the Query Service integration's dispatch
+mechanism (BYOC-DESIGN.md's AlwaysOn/Arrow vs. code-interpreter branch
+fork) is still unresolved. This contract shape (one row of scalar
+arguments in, one scalar value out) matches a synchronous per-call SQL
+UDF invocation and may need revision once the branch decision lands.
+"""
+
+from typing import (
+    Any,
+    List,
+    Optional,
+)
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+)
+
+
+class QueryServiceUdfV1Argument(BaseModel):
+    """A single positional argument passed into the UDF for one row"""
+
+    name: Optional[str] = Field(
+        default=None,
+        description="Parameter name as declared in the UDF signature",
+        examples=["amount"],
+    )
+    value: Optional[Any] = Field(
+        default=None,
+        description="Argument value for this row. Type follows the UDF's declared parameter type.",
+        examples=[42.5],
+    )
+    model_config = ConfigDict(extra="ignore")
+
+
+class QueryServiceUdfV1Row(BaseModel):
+    """One row of arguments to evaluate the UDF against"""
+
+    arguments: List[QueryServiceUdfV1Argument] = Field(
+        default_factory=list,
+        description="Ordered arguments for this row, matching the UDF signature",
+    )
+    model_config = ConfigDict(extra="ignore")
+
+
+class QueryServiceUdfV1Result(BaseModel):
+    """Evaluated result for one input row"""
+
+    value: Optional[Any] = Field(
+        default=None,
+        description="Return value for this row. Type follows the UDF's declared return type.",
+        examples=[127.5],
+    )
+    model_config = ConfigDict(extra="ignore")
+
+
+class QueryServiceUdfV1Request(BaseModel):
+    """Request for Query Service UDF invocation"""
+
+    rows: List[QueryServiceUdfV1Row] = Field(
+        default_factory=list, description="Rows to evaluate the UDF against"
+    )
+    model_config = ConfigDict(extra="ignore")
+
+
+class QueryServiceUdfV1Response(BaseModel):
+    """Response for Query Service UDF invocation"""
+
+    results: List[QueryServiceUdfV1Result] = Field(
+        default_factory=list, description="Evaluated results, one per input row"
+    )