Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions src/datacustomcode/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
90 changes: 90 additions & 0 deletions src/datacustomcode/function/feature_types/query_service_udf.py
Original file line number Diff line number Diff line change
@@ -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"
)