Merge pull request #94 from forcedotcom/chunking_contract

Search Index Chunking contract

Author: ritaagarwala-sf

.github/workflows/sf_cli_integration.yml

@@ -200,6 +200,10 @@ jobs:
           echo "::error::testFunction/.datacustomcode_proj/sdk_config.json not found after function init."
           exit 1
         }
+        test -f testFunction/payload/tests/test.json || {
+          echo "::error::testFunction/payload/tests/test.json not found after function init."
+          exit 1
+        }
 
       # ── Function: scan ────────────────────────────────────────────────────────
 
@@ -251,14 +255,14 @@ jobs:
 
       # ── Function: run ─────────────────────────────────────────────────────────
 
-    - name: '[function] run — sf data-code-extension function run --entrypoint testFunction/payload/entrypoint.py -o dev1'
+    - name: '[function] run — sf data-code-extension function run --entrypoint testFunction/payload/entrypoint.py  --test-with testFunction/payload/tests/test.json -o dev1'
       run: |
         sf data-code-extension function run \
           --entrypoint testFunction/payload/entrypoint.py \
-          -o dev1 || {
-          echo "::error::sf data-code-extension function run FAILED. Check mock server output above; the --entrypoint flag or SF CLI org auth contract may have changed."
-          exit 1
-        }
+          --test-with testFunction/payload/tests/test.json || {
+            echo "::error::sf data-code-extension function run FAILED. Check mock server output above; the --entrypoint flag or SF CLI org auth contract may have changed."
+            exit 1
+          }
 
       # ── Function: deploy ─────────────────────────────────────────────────────
 
@@ -270,7 +274,6 @@ jobs:
           --description "Test function deploy" \
           --package-dir testFunction/payload \
           --cpu-size CPU_2XL \
-          --function-invoke-opt UnstructuredChunking \
           -o dev1 || {
           echo "::error::sf data-code-extension function deploy FAILED. Check mock server output above for which endpoint failed. The deploy command flags or API contract may have changed."
           exit 1

src/datacustomcode/cli.py

@@ -27,6 +27,13 @@
 
 from datacustomcode import AuthType
 from datacustomcode.auth import configure_oauth_tokens
+from datacustomcode.constants import (
+    CONFIG_FILE,
+    ENTRYPOINT_FILE,
+    PAYLOAD_DIR,
+    TEST_FILE,
+    TESTS_DIR,
+)
 from datacustomcode.scan import find_base_directory, get_package_type
 
 
@@ -74,6 +81,30 @@ def _configure_client_credentials(
     )
 
 
+def _generate_function_test_file(entrypoint_path: str) -> Optional[str]:
+    """Generate test.json file for a function.
+
+    Args:
+        entrypoint_path: Path to the function's entrypoint.py
+
+    Returns:
+        Path to generated test.json, or None if generation failed
+    """
+    from datacustomcode.function_utils import generate_test_json
+
+    tests_dir = os.path.join(os.path.dirname(entrypoint_path), TESTS_DIR)
+    os.makedirs(tests_dir, exist_ok=True)
+    test_json_path = os.path.join(tests_dir, TEST_FILE)
+
+    try:
+        generate_test_json(entrypoint_path, test_json_path)
+        logger.debug(f"Generated test JSON at {test_json_path}")
+        return test_json_path
+    except Exception as e:
+        logger.warning(f"Could not generate test.json: {e}")
+        return None
+
+
 @cli.command()
 @click.option("--profile", default="default", help="Credential profile name")
 @click.option(
@@ -162,7 +193,6 @@ def zip(path: str, network: str):
 
     Choose based on your workload requirements.""",
 )
-@click.option("--function-invoke-opt")
 @click.option(
     "--sf-cli-org",
     default=None,
@@ -176,13 +206,14 @@ def deploy(
     cpu_size: str,
     profile: str,
     network: str,
-    function_invoke_opt: str,
     sf_cli_org: Optional[str],
 ):
+    from datacustomcode.constants import USE_IN_FEATURE_MAPPING_FOR_CONNECT_API
     from datacustomcode.deploy import (
         COMPUTE_TYPES,
         CodeExtensionMetadata,
         deploy_full,
+        infer_use_in_feature,
     )
     from datacustomcode.token_provider import (
         CredentialsTokenProvider,
@@ -211,15 +242,24 @@ def deploy(
     )
 
     if package_type == "function":
-        if not function_invoke_opt:
+        # Infer use_in_feature from function signature
+        entrypoint_path = os.path.join(path, ENTRYPOINT_FILE)
+        use_in_feature = infer_use_in_feature(entrypoint_path)
+        if use_in_feature:
+            logger.info(f"Inferred use_in_feature: {use_in_feature}")
+        else:
             click.secho(
-                "Error: Function invoke options are required for function package type",
+                "Error: Could not infer function invoke options. "
+                "Please provide --use-in-feature",
                 fg="red",
             )
             raise click.Abort()
-        else:
-            function_invoke_options = function_invoke_opt.split(",")
-            metadata.functionInvokeOptions = function_invoke_options
+
+        # Map user-provided feature names to API names
+        mapped_feature = USE_IN_FEATURE_MAPPING_FOR_CONNECT_API.get(
+            use_in_feature, use_in_feature
+        )
+        metadata.functionInvokeOptions = [mapped_feature]
 
     try:
         if sf_cli_org:
@@ -238,7 +278,12 @@ def deploy(
 @click.option(
     "--code-type", default="script", type=click.Choice(["script", "function"])
 )
-def init(directory: str, code_type: str):
+@click.option(
+    "--use-in-feature",
+    default="SearchIndexChunking",
+    help="Feature where this function will be used (only applicable for function).",
+)
+def init(directory: str, code_type: str, use_in_feature: Optional[str]):
     from datacustomcode.scan import (
         dc_config_json_from_file,
         update_config,
@@ -250,9 +295,9 @@ def init(directory: str, code_type: str):
     if code_type == "script":
         copy_script_template(directory)
     elif code_type == "function":
-        copy_function_template(directory)
-    entrypoint_path = os.path.join(directory, "payload", "entrypoint.py")
-    config_location = os.path.join(os.path.dirname(entrypoint_path), "config.json")
+        copy_function_template(directory, use_in_feature)
+    entrypoint_path = os.path.join(directory, PAYLOAD_DIR, ENTRYPOINT_FILE)
+    config_location = os.path.join(os.path.dirname(entrypoint_path), CONFIG_FILE)
 
     # Write package type to SDK-specific config
     sdk_config = {"type": code_type}
@@ -265,6 +310,7 @@ def init(directory: str, code_type: str):
     updated_config_json = update_config(entrypoint_path)
     with open(config_location, "w") as f:
         json.dump(updated_config_json, f, indent=2)
+
     click.echo(
         "Start developing by updating the code in "
         + click.style(entrypoint_path, fg="blue", bold=True)
@@ -275,6 +321,24 @@ def init(directory: str, code_type: str):
         + " to automatically update config.json when you make changes to your code"
     )
 
+    # Generate test.json for functions
+    if code_type == "function":
+        test_json_path = _generate_function_test_file(entrypoint_path)
+        if test_json_path:
+            click.echo(
+                "Generated test file at "
+                + click.style(test_json_path, fg="blue", bold=True)
+            )
+            click.echo(
+                "Test your function locally with "
+                + click.style(
+                    f"datacustomcode run {entrypoint_path} "
+                    f"--test-with {test_json_path}",
+                    fg="blue",
+                    bold=True,
+                )
+            )
+
 
 @cli.command()
 @click.argument("filename")
@@ -286,7 +350,7 @@ def init(directory: str, code_type: str):
 def scan(filename: str, config: str, dry_run: bool, no_requirements: bool):
     from datacustomcode.scan import update_config, write_requirements_file
 
-    config_location = config or os.path.join(os.path.dirname(filename), "config.json")
+    config_location = config or os.path.join(os.path.dirname(filename), CONFIG_FILE)
     click.echo(
         "Dumping scan results to config file: "
         + click.style(config_location, fg="blue", bold=True)
@@ -312,6 +376,12 @@ def scan(filename: str, config: str, dry_run: bool, no_requirements: bool):
 @click.option("--config-file", default=None)
 @click.option("--dependencies", default=[], multiple=True)
 @click.option("--profile", default="default")
+@click.option(
+    "--test-with",
+    default=None,
+    type=click.Path(exists=True),
+    help="Path to test JSON file for function testing",
+)
 @click.option(
     "--sf-cli-org",
     default=None,
@@ -322,10 +392,16 @@ def run(
     config_file: Union[str, None],
     dependencies: List[str],
     profile: str,
+    test_with: Optional[str],
     sf_cli_org: Optional[str],
 ):
     from datacustomcode.run import run_entrypoint
 
     run_entrypoint(
-        entrypoint, config_file, dependencies, profile, sf_cli_org=sf_cli_org
+        entrypoint,
+        config_file,
+        dependencies,
+        profile,
+        test_file=test_with,
+        sf_cli_org=sf_cli_org,
     )

src/datacustomcode/constants.py

@@ -0,0 +1,45 @@
+# 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.
+
+"""Constants used throughout the datacustomcode package."""
+
+# File and directory names
+ENTRYPOINT_FILE = "entrypoint.py"
+CONFIG_FILE = "config.json"
+PAYLOAD_DIR = "payload"
+TESTS_DIR = "tests"
+TEST_FILE = "test.json"
+REQUIREMENTS_FILE = "requirements.txt"
+
+# Default values
+DEFAULT_PROFILE = "default"
+DEFAULT_NETWORK = "default"
+DEFAULT_CPU_SIZE = "CPU_2XL"
+
+# Feature to template folder mapping
+FEATURE_TEMPLATE_MAPPING = {
+    "SearchIndexChunking": "chunking",
+}
+
+# Feature name to Connect API name mapping
+USE_IN_FEATURE_MAPPING_FOR_CONNECT_API = {
+    "SearchIndexChunking": "UnstructuredChunking",
+}
+
+# Pydantic request/response type names to feature names
+REQUEST_TYPE_TO_FEATURE = {
+    "SearchIndexChunkingV1Request": "SearchIndexChunking",
+    "SearchIndexChunkingV1Response": "SearchIndexChunking",
+}

src/datacustomcode/deploy.py

@@ -35,6 +35,7 @@
 import requests
 
 from datacustomcode.cmd import cmd_output
+from datacustomcode.constants import REQUEST_TYPE_TO_FEATURE
 from datacustomcode.scan import find_base_directory, get_package_type
 
 DATA_CUSTOM_CODE_PATH = "services/data/v63.0/ssot/data-custom-code"
@@ -65,6 +66,40 @@ def _sanitize_api_name(name: str) -> str:
     return sanitized
 
 
+def infer_use_in_feature(entrypoint_path: str) -> Union[str, None]:
+    """Infer the use_in_feature from function signature.
+
+    Checks both the request parameter type and return type annotation.
+    Both must map to the same feature for a valid inference.
+
+    Uses static AST parsing to avoid importing dependencies.
+
+    Args:
+        entrypoint_path: Path to the entrypoint.py file
+
+    Returns:
+        The feature name if both request and response match, None otherwise
+    """
+    from datacustomcode.function_utils import inspect_function_types_static
+
+    request_type_name, response_type_name = inspect_function_types_static(
+        entrypoint_path
+    )
+
+    if not request_type_name or not response_type_name:
+        return None
+
+    # Look up features for both types
+    request_feature = REQUEST_TYPE_TO_FEATURE.get(request_type_name)
+    response_feature = REQUEST_TYPE_TO_FEATURE.get(response_type_name)
+
+    # Both must be present and must match
+    if request_feature and response_feature and request_feature == response_feature:
+        return request_feature
+
+    return None
+
+
 class CodeExtensionMetadata(BaseModel):
     name: str
     version: str