feat: support sb settlement client (#1763)

* initial worker changes for supporting sb settlement client

* feedback

* backport to v2, add unit tests

* update

* install into path

* no deps

* does this work

* there is potential

* rewrite tests

* clean up

* sb in its own dir

* fix test

* fix test

* fix test

* fix 313

* change return empty string to return none

* change return empty string to return none pt2

* fix tests

* move 313 tests

Author: hallvictoria

eng/scripts/install-dependencies.sh

@@ -11,4 +11,8 @@ if [[ $1 != "3.7" ]]; then
 fi
 if [[ $1 != "3.7" && $1 != "3.8" ]]; then
     python -m pip install --pre -U -e $2/[test-deferred-bindings]
+
+    SERVICEBUS_DIR="./servicebus_dir"
+    python -m pip install --pre -U --target "$SERVICEBUS_DIR" azurefunctions-extensions-bindings-servicebus==1.0.0b2
+    python -c "import sys; sys.path.insert(0, '$SERVICEBUS_DIR'); import azurefunctions.extensions.bindings.servicebus as sb; print('servicebus version:', sb.__version__)"
 fi

eng/templates/jobs/ci-emulator-tests.yml

@@ -108,6 +108,7 @@ jobs:
           AzureWebJobsSQLPassword: $(AzureWebJobsSQLPassword)
         displayName: 'Install Azurite and Start ServiceBus Emulator'
       - bash: |
+            export PYTHONPATH="$(Build.SourcesDirectory)/servicebus_dir:$PYTHONPATH"
             python -m pytest -q --dist loadfile --reruns 4 tests/emulator_tests/test_servicebus_functions.py
         env:
           AzureWebJobsStorage: $(AzureWebJobsStorage)

runtimes/v2/azure_functions_runtime/bindings/meta.py

@@ -15,6 +15,7 @@
     CUSTOMER_PACKAGES_PATH,
     HTTP,
     HTTP_TRIGGER,
+    SERVICE_BUS_CLIENT_NAME
 )
 from ..utils.helpers import set_sdk_version
 
@@ -286,3 +287,39 @@ def get_deferred_raw_bindings(indexed_function, input_types):
     raw_bindings, bindings_logs = DEFERRED_BINDING_REGISTRY.get_raw_bindings(
         indexed_function, input_types)
     return raw_bindings, bindings_logs
+
+
+def get_settlement_client():
+    return DEFERRED_BINDING_REGISTRY.get(SERVICE_BUS_CLIENT_NAME).get_client()
+
+
+def validate_settlement_param(params: dict,
+                              bound_params: dict,
+                              annotations: dict) -> str:
+    """
+    Checks if the settlement client is enabled for a given function.
+
+    If there is more than one param that is not bound, return an empty string
+    indicating no settlement client support. This is a bad app.
+
+    If there is only one unbound param, check if it's a type that is supported
+    by the settlement client. If so, return the param_name, where param_name
+    is the name of the param that is supported. If not, return an empty string
+    indicating no settlement client support.
+
+    Note: If a param does not have a type annotation, it will be skipped and not
+    considered for settlement client support.
+    """
+    if len(set(params) - set(bound_params)) > 1:
+        return None
+
+    # There is only one unbound param, check the type
+    settlement_param = next(iter(set(params) - set(bound_params)))
+    try:
+        param_type = annotations.get(settlement_param)
+        # Check if the type is a supported type for the settlement client
+        if DEFERRED_BINDING_REGISTRY.check_supported_grpc_client_type(param_type):
+            return settlement_param
+    except Exception:
+        param_type = None
+    return None

runtimes/v2/azure_functions_runtime/functions.py

@@ -11,7 +11,8 @@
 from .bindings.meta import (has_implicit_output,
                             check_deferred_bindings_enabled,
                             check_output_type_annotation,
-                            check_input_type_annotation)
+                            check_input_type_annotation,
+                            validate_settlement_param)
 from .utils.constants import HTTP_TRIGGER
 from .utils.typing_inspect import is_generic_type, get_origin, get_args  # type: ignore
 
@@ -33,6 +34,7 @@ class FunctionInfo(typing.NamedTuple):
     has_return: bool
     is_http_func: bool
     deferred_bindings_enabled: bool
+    settlement_client_arg: typing.Optional[str]
 
     input_types: typing.Mapping[str, ParamTypeInfo]
     output_types: typing.Mapping[str, ParamTypeInfo]
@@ -138,12 +140,20 @@ def validate_function_params(params: dict, bound_params: dict,
                                  protos):
         logger.debug("Params: %s, BoundParams: %s, Annotations: %s, FuncName: %s",
                      params, bound_params, annotations, func_name)
+        settlement_client_arg = None
         if set(params) - set(bound_params):
-            raise FunctionLoadError(
-                func_name,
-                'Function parameter mismatch — the following trigger/input bindings '
-                'are declared in Python but missing from the '
-                'function decorator: ' + repr(set(params) - set(bound_params)))
+            # Check for settlement client support for the missing parameters
+            settlement_client_arg = validate_settlement_param(
+                params, bound_params, annotations)
+            if settlement_client_arg is not None:
+                params.pop(settlement_client_arg)
+            else:
+                # Not supported by settlement client, raise error for missing parameters
+                raise FunctionLoadError(
+                    func_name,
+                    'the following parameters are declared in Python '
+                    'but not in the function definition (function.json or '
+                    f'function decorators):  {set(params) - set(bound_params)!r}')
 
         if set(bound_params) - set(params):
             raise FunctionLoadError(
@@ -278,7 +288,8 @@ def validate_function_params(params: dict, bound_params: dict,
                 output_types[param.name] = param_type_info
             else:
                 input_types[param.name] = param_type_info
-        return input_types, output_types, fx_deferred_bindings_enabled
+        return (input_types, output_types, fx_deferred_bindings_enabled,
+                settlement_client_arg)
 
     @staticmethod
     def get_function_return_type(annotations: dict, has_explicit_return: bool,
@@ -330,6 +341,7 @@ def add_func_to_registry_and_return_funcinfo(
             has_explicit_return: bool,
             has_implicit_return: bool,
             deferred_bindings_enabled: bool,
+            settlement_client_arg: typing.Optional[str],
             input_types: typing.Dict[str, ParamTypeInfo],
             output_types: typing.Dict[str, ParamTypeInfo],
             return_type: str):
@@ -355,6 +367,7 @@ def add_func_to_registry_and_return_funcinfo(
             has_return=has_explicit_return or has_implicit_return,
             is_http_func=is_http_func,
             deferred_bindings_enabled=deferred_bindings_enabled,
+            settlement_client_arg=settlement_client_arg,
             input_types=input_types,
             output_types=output_types,
             return_type=return_type,
@@ -412,7 +425,8 @@ def add_indexed_function(self, function, protos):
                                                     func_name)
 
         (input_types, output_types,
-         deferred_bindings_enabled) = self.validate_function_params(
+         deferred_bindings_enabled,
+         settlement_client_arg) = self.validate_function_params(
             params,
             bound_params,
             annotations,
@@ -431,5 +445,6 @@ def add_indexed_function(self, function, protos):
                 func, func_name, function_id, func_dir,
                 requires_context, has_explicit_return,
                 has_implicit_return, deferred_bindings_enabled,
+                settlement_client_arg,
                 input_types, output_types,
                 return_type)

runtimes/v2/azure_functions_runtime/handle_event.py

@@ -22,6 +22,7 @@
 
 from .bindings.context import get_context
 from .bindings.meta import (from_incoming_proto,
+                            get_settlement_client,
                             is_trigger_binding,
                             load_binding_registry,
                             to_outgoing_param_binding,
@@ -224,6 +225,9 @@ async def invocation_request(request):
         if fi.requires_context:
             args['context'] = fi_context
 
+        if fi.settlement_client_arg is not None:
+            args[fi.settlement_client_arg] = get_settlement_client()
+
         if fi.output_types:
             for name in fi.output_types:
                 args[name] = Out()

runtimes/v2/azure_functions_runtime/utils/constants.py

@@ -11,6 +11,7 @@
 MODULE_NOT_FOUND_TS_URL = "https://aka.ms/functions-modulenotfound"
 PYTHON_LANGUAGE_RUNTIME = "python"
 RETRY_POLICY = "retry_policy"
+SERVICE_BUS_CLIENT_NAME = "serviceBusClient"
 TRUE = "true"
 TRACEPARENT = "traceparent"
 TRACESTATE = "tracestate"

runtimes/v2/tests/unittests/test_deferred_bindings.py

@@ -16,11 +16,15 @@
                                                      ContainerClient,
                                                      StorageStreamDownloader)
 from azurefunctions.extensions.bindings.eventhub import EventData, EventDataConverter
+from azurefunctions.extensions.base import GrpcClientType
 
 EVENTHUB_SAMPLE_CONTENT = b"\x00Sr\xc1\x8e\x08\xa3\x1bx-opt-sequence-number-epochT\xff\xa3\x15x-opt-sequence-numberU\x04\xa3\x0cx-opt-offset\x81\x00\x00\x00\x01\x00\x00\x010\xa3\x13x-opt-enqueued-time\x00\xa3\x1dcom.microsoft:datetime-offset\x81\x08\xddW\x05\xc3Q\xcf\x10\x00St\xc1I\x02\xa1\rDiagnostic-Id\xa1700-bdc3fde4889b4e907e0c9dcb46ff8d92-21f637af293ef13b-00\x00Su\xa0\x08message1"  # noqa: E501
 
 
 class TestDeferredBindingsEnabled(testutils.AsyncTestCase):
+    def setUp(self):
+        # Initialize DEFERRED_BINDING_REGISTRY
+        meta.load_binding_registry()
 
     @unittest.skip("TODO: Move to emulator.")
     def test_mbd_deferred_bindings_enabled_decode(self):
@@ -98,3 +102,46 @@ async def test_check_deferred_bindings_enabled(self):
             ContainerClient, True), (True, True))
         self.assertEqual(meta.check_deferred_bindings_enabled(
             StorageStreamDownloader, True), (True, True))
+
+    async def test_valid_settlement_param(self):
+        params = {'param1', 'param2', 'param3'}
+        bound_params = {'param1', 'param2'}
+        annotations = {
+            'param1': func.InputStream,
+            'param2': func.Out[str],
+            'param3': GrpcClientType
+        }
+
+        settlement_client_arg = meta.validate_settlement_param(
+            params, bound_params, annotations)
+
+        self.assertEqual(settlement_client_arg, 'param3')
+
+    async def test_invalid_settlement_param(self):
+        params = {'param1', 'param2', 'param3'}
+        bound_params = {'param1', 'param2'}
+        annotations = {
+            'param1': func.InputStream,
+            'param2': func.Out[str],
+            'param3': str
+        }
+
+        settlement_client_arg = meta.validate_settlement_param(
+            params, bound_params, annotations)
+
+        self.assertEqual(settlement_client_arg, None)
+
+    async def test_invalid_settlement_param_multiple(self):
+        params = {'param1', 'param2', 'param3', 'param4'}
+        bound_params = {'param1', 'param2'}
+        annotations = {
+            'param1': func.InputStream,
+            'param2': func.Out[str],
+            'param3': GrpcClientType,
+            'param4': str
+        }
+
+        settlement_client_arg = meta.validate_settlement_param(
+            params, bound_params, annotations)
+
+        self.assertEqual(settlement_client_arg, None)

workers/azure_functions_worker/bindings/__init__.py

@@ -4,6 +4,7 @@
 from .tracecontext import TraceContext  # isort: skip
 from .context import Context
 from .meta import (
+    get_settlement_client,
     check_deferred_bindings_enabled,
     check_input_type_annotation,
     check_output_type_annotation,
@@ -12,18 +13,20 @@
     has_implicit_output,
     is_trigger_binding,
     load_binding_registry,
+    validate_settlement_param,
     to_outgoing_param_binding,
     to_outgoing_proto,
 )
 from .out import Out
 
 __all__ = (
     'Out', 'Context',
+    'get_settlement_client',
     'is_trigger_binding',
     'load_binding_registry',
     'check_input_type_annotation', 'check_output_type_annotation',
     'has_implicit_output',
     'from_incoming_proto', 'to_outgoing_proto', 'TraceContext', 'RetryContext',
     'to_outgoing_param_binding', 'check_deferred_bindings_enabled',
-    'get_deferred_raw_bindings'
+    'get_deferred_raw_bindings', 'validate_settlement_param'
 )

workers/azure_functions_worker/bindings/meta.py

@@ -10,6 +10,7 @@
     CUSTOMER_PACKAGES_PATH,
     HTTP,
     HTTP_TRIGGER,
+    SERVICE_BUS_CLIENT_NAME
 )
 from ..http_v2 import HttpV2Registry
 from ..logging import logger
@@ -288,9 +289,9 @@ def deferred_bindings_decode(binding: typing.Any,
     return binding.decode(datum, trigger_metadata=metadata, pytype=pytype)
 
 
-def check_deferred_bindings_enabled(param_anno: type,
-                                    deferred_bindings_enabled: bool) -> (bool,
-                                                                         bool):
+def check_deferred_bindings_enabled(
+        param_anno: type,
+        deferred_bindings_enabled: bool) -> typing.Tuple[bool, bool]:
     """
     Checks if deferred bindings is enabled at fx and single binding level
 
@@ -314,3 +315,39 @@ def get_deferred_raw_bindings(indexed_function, input_types):
     raw_bindings, bindings_logs = DEFERRED_BINDING_REGISTRY.get_raw_bindings(
         indexed_function, input_types)
     return raw_bindings, bindings_logs
+
+
+def get_settlement_client():
+    return DEFERRED_BINDING_REGISTRY.get(SERVICE_BUS_CLIENT_NAME).get_client()
+
+
+def validate_settlement_param(params: dict,
+                              bound_params: dict,
+                              annotations: dict) -> str:
+    """
+    Checks if the settlement client is enabled for a given function.
+
+    If there is more than one param that is not bound, return an empty string
+    indicating no settlement client support. This is a bad app.
+
+    If there is only one unbound param, check if it's a type that is supported
+    by the settlement client. If so, return the param_name, where param_name
+    is the name of the param that is supported. If not, return an empty string
+    indicating no settlement client support.
+
+    Note: If a param does not have a type annotation, it will be skipped and not
+    considered for settlement client support.
+    """
+    if len(set(params) - set(bound_params)) > 1:
+        return None
+
+    # There is only one unbound param, check the type
+    settlement_param = next(iter(set(params) - set(bound_params)))
+    try:
+        param_type = annotations.get(settlement_param)
+        # Check if the type is a supported type for the settlement client
+        if DEFERRED_BINDING_REGISTRY.check_supported_grpc_client_type(param_type):
+            return settlement_param
+    except Exception:
+        param_type = None
+    return None

workers/azure_functions_worker/constants.py

@@ -115,3 +115,6 @@
     '3.11': '2027-04',
     '3.12': '2028-04'
 }
+
+# Service Bus Client Name
+SERVICE_BUS_CLIENT_NAME = "serviceBusClient"