diff --git a/.coveragerc-py37 b/.coveragerc-py37 index b1c98d23..efb63fec 100644 --- a/.coveragerc-py37 +++ b/.coveragerc-py37 @@ -12,14 +12,11 @@ omit = [report] exclude_lines = - # Have to re-enable the standard pragma pragma: no cover - - # Don't complain about async-specific imports and code from functions_framework.aio import from functions_framework._http.asgi import from functions_framework._http.gunicorn import UvicornApplication - - # Exclude async-specific classes and functions in execution_id.py class AsgiMiddleware: - def set_execution_context_async \ No newline at end of file + def set_execution_context_async + return create_asgi_app_from_module + app = create_asgi_app\(target, source, signature_type\) \ No newline at end of file diff --git a/src/functions_framework/__init__.py b/src/functions_framework/__init__.py index 22fbf44c..31169f4c 100644 --- a/src/functions_framework/__init__.py +++ b/src/functions_framework/__init__.py @@ -327,6 +327,16 @@ def crash_handler(e): def create_app(target=None, source=None, signature_type=None): + """Create an app for the function. + + Args: + target: The name of the target function to invoke + source: The source file containing the function + signature_type: The signature type of the function + + Returns: + A Flask WSGI app or Starlette ASGI app depending on function decorators + """ target = _function_registry.get_function_target(target) source = _function_registry.get_function_source(source) @@ -370,6 +380,7 @@ def handle_none(rv): setup_logging() _app.wsgi_app = execution_id.WsgiMiddleware(_app.wsgi_app) + # Execute the module, within the application context with _app.app_context(): try: @@ -394,6 +405,23 @@ def function(*_args, **_kwargs): # command fails. raise e from None + use_asgi = target in _function_registry.ASGI_FUNCTIONS + if use_asgi: + # This function needs ASGI, delegate to create_asgi_app + # Note: @aio decorators only register functions in ASGI_FUNCTIONS when the + # module is imported. We can't know if a function uses @aio until after + # we load the module. + # + # To avoid loading modules twice, we always create a Flask app first, load the + # module within its context, then check if ASGI is needed. This results in an + # unused Flask app for ASGI functions, but we accept this memory overhead as a + # trade-off. + from functions_framework.aio import create_asgi_app_from_module + + return create_asgi_app_from_module( + target, source, signature_type, source_module, spec + ) + # Get the configured function signature type signature_type = _function_registry.get_func_signature_type(target, signature_type) diff --git a/src/functions_framework/_cli.py b/src/functions_framework/_cli.py index c2ba9f4b..48455ea6 100644 --- a/src/functions_framework/_cli.py +++ b/src/functions_framework/_cli.py @@ -16,7 +16,7 @@ import click -from functions_framework import create_app +from functions_framework import _function_registry, create_app from functions_framework._http import create_server @@ -39,11 +39,10 @@ help="Use ASGI server for function execution", ) def _cli(target, source, signature_type, host, port, debug, asgi): - if asgi: # pragma: no cover + if asgi: from functions_framework.aio import create_asgi_app app = create_asgi_app(target, source, signature_type) else: app = create_app(target, source, signature_type) - create_server(app, debug).run(host, port) diff --git a/src/functions_framework/_function_registry.py b/src/functions_framework/_function_registry.py index 2214b5fd..1f08c794 100644 --- a/src/functions_framework/_function_registry.py +++ b/src/functions_framework/_function_registry.py @@ -40,6 +40,10 @@ # Keys are the user function name, values are the type of the function input INPUT_TYPE_MAP = {} +# ASGI_FUNCTIONS stores function names that require ASGI mode. +# Functions decorated with @aio.http or @aio.cloud_event are added here. +ASGI_FUNCTIONS = set() + def get_user_function(source, source_module, target): """Returns user function, raises exception for invalid function.""" diff --git a/src/functions_framework/aio/__init__.py b/src/functions_framework/aio/__init__.py index 8e5f9dc7..a56fe942 100644 --- a/src/functions_framework/aio/__init__.py +++ b/src/functions_framework/aio/__init__.py @@ -62,6 +62,7 @@ def cloud_event(func: CloudEventFunction) -> CloudEventFunction: _function_registry.REGISTRY_MAP[func.__name__] = ( _function_registry.CLOUDEVENT_SIGNATURE_TYPE ) + _function_registry.ASGI_FUNCTIONS.add(func.__name__) if inspect.iscoroutinefunction(func): @functools.wraps(func) @@ -82,6 +83,7 @@ def http(func: HTTPFunction) -> HTTPFunction: _function_registry.REGISTRY_MAP[func.__name__] = ( _function_registry.HTTP_SIGNATURE_TYPE ) + _function_registry.ASGI_FUNCTIONS.add(func.__name__) if inspect.iscoroutinefunction(func): @@ -213,6 +215,29 @@ async def __call__(self, scope, receive, send): # Don't re-raise to prevent starlette from printing traceback again +def create_asgi_app_from_module(target, source, signature_type, source_module, spec): + """Create an ASGI application from an already-loaded module. + + Args: + target: The name of the target function to invoke + source: The source file containing the function + signature_type: The signature type of the function + source_module: The already-loaded module + spec: The module spec + + Returns: + A Starlette ASGI application instance + """ + enable_id_logging = _enable_execution_id_logging() + if enable_id_logging: # pragma: no cover + _configure_app_execution_id_logging() + + function = _function_registry.get_user_function(source, source_module, target) + signature_type = _function_registry.get_func_signature_type(target, signature_type) + + return _create_asgi_app_with_function(function, signature_type, enable_id_logging) + + def create_asgi_app(target=None, source=None, signature_type=None): """Create an ASGI application for the function. @@ -243,6 +268,11 @@ def create_asgi_app(target=None, source=None, signature_type=None): function = _function_registry.get_user_function(source, source_module, target) signature_type = _function_registry.get_func_signature_type(target, signature_type) + return _create_asgi_app_with_function(function, signature_type, enable_id_logging) + + +def _create_asgi_app_with_function(function, signature_type, enable_id_logging): + """Create an ASGI app with the given function and signature type.""" is_async = inspect.iscoroutinefunction(function) routes = [] if signature_type == _function_registry.HTTP_SIGNATURE_TYPE: diff --git a/tests/test_cli.py b/tests/test_cli.py index 4e5a0a08..75c93f20 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os +import pathlib import sys import pretend @@ -20,9 +22,30 @@ from click.testing import CliRunner import functions_framework +import functions_framework._function_registry as _function_registry from functions_framework._cli import _cli +# Conditional import for Starlette (Python 3.8+) +if sys.version_info >= (3, 8): + from starlette.applications import Starlette +else: + Starlette = None + + +@pytest.fixture(autouse=True) +def clean_registries(): + """Clean up both REGISTRY_MAP and ASGI_FUNCTIONS registries.""" + original_registry_map = _function_registry.REGISTRY_MAP.copy() + original_asgi = _function_registry.ASGI_FUNCTIONS.copy() + _function_registry.REGISTRY_MAP.clear() + _function_registry.ASGI_FUNCTIONS.clear() + yield + _function_registry.REGISTRY_MAP.clear() + _function_registry.REGISTRY_MAP.update(original_registry_map) + _function_registry.ASGI_FUNCTIONS.clear() + _function_registry.ASGI_FUNCTIONS.update(original_asgi) + def test_cli_no_arguments(): runner = CliRunner() @@ -124,3 +147,19 @@ def test_asgi_cli(monkeypatch): assert result.exit_code == 0 assert create_asgi_app.calls == [pretend.call("foo", None, "http")] assert asgi_server.run.calls == [pretend.call("0.0.0.0", 8080)] + + +def test_cli_auto_detects_asgi_decorator(): + """Test that CLI auto-detects @aio decorated functions without --asgi flag.""" + # Use the actual async_decorator.py test file which has @aio.http decorated functions + test_functions_dir = pathlib.Path(__file__).parent / "test_functions" / "decorators" + source = test_functions_dir / "async_decorator.py" + + # Call create_app without any asgi flag - should auto-detect + app = functions_framework.create_app(target="function_http", source=str(source)) + + # Verify it created a Starlette app (ASGI) + assert isinstance(app, Starlette) + + # Verify the function was registered in ASGI_FUNCTIONS + assert "function_http" in _function_registry.ASGI_FUNCTIONS diff --git a/tests/test_decorator_functions.py b/tests/test_decorator_functions.py index 435aa815..3a6e5e99 100644 --- a/tests/test_decorator_functions.py +++ b/tests/test_decorator_functions.py @@ -19,6 +19,8 @@ from cloudevents import conversion as ce_conversion from cloudevents.http import CloudEvent +import functions_framework._function_registry as registry + # Conditional import for Starlette if sys.version_info >= (3, 8): from starlette.testclient import TestClient as StarletteTestClient @@ -35,6 +37,21 @@ TEST_FUNCTIONS_DIR = pathlib.Path(__file__).resolve().parent / "test_functions" + +@pytest.fixture(autouse=True) +def clean_registries(): + """Clean up both REGISTRY_MAP and ASGI_FUNCTIONS registries.""" + original_registry_map = registry.REGISTRY_MAP.copy() + original_asgi = registry.ASGI_FUNCTIONS.copy() + registry.REGISTRY_MAP.clear() + registry.ASGI_FUNCTIONS.clear() + yield + registry.REGISTRY_MAP.clear() + registry.REGISTRY_MAP.update(original_registry_map) + registry.ASGI_FUNCTIONS.clear() + registry.ASGI_FUNCTIONS.update(original_asgi) + + # Python 3.5: ModuleNotFoundError does not exist try: _ModuleNotFoundError = ModuleNotFoundError @@ -128,3 +145,33 @@ def test_aio_http_dict_response(): resp = client.post("/") assert resp.status_code == 200 assert resp.json() == {"message": "hello", "count": 42, "success": True} + + +def test_aio_decorators_register_asgi_functions(): + """Test that @aio decorators add function names to ASGI_FUNCTIONS registry.""" + from functions_framework.aio import cloud_event, http + + @http + async def test_http_func(request): + return "test" + + @cloud_event + async def test_cloud_event_func(event): + pass + + assert "test_http_func" in registry.ASGI_FUNCTIONS + assert "test_cloud_event_func" in registry.ASGI_FUNCTIONS + + assert registry.REGISTRY_MAP["test_http_func"] == "http" + assert registry.REGISTRY_MAP["test_cloud_event_func"] == "cloudevent" + + @http + def test_http_sync(request): + return "sync" + + @cloud_event + def test_cloud_event_sync(event): + pass + + assert "test_http_sync" in registry.ASGI_FUNCTIONS + assert "test_cloud_event_sync" in registry.ASGI_FUNCTIONS diff --git a/tests/test_function_registry.py b/tests/test_function_registry.py index e3ae3c7e..5b517cdc 100644 --- a/tests/test_function_registry.py +++ b/tests/test_function_registry.py @@ -13,9 +13,25 @@ # limitations under the License. import os +import pytest + from functions_framework import _function_registry +@pytest.fixture(autouse=True) +def clean_registries(): + """Clean up both REGISTRY_MAP and ASGI_FUNCTIONS registries.""" + original_registry_map = _function_registry.REGISTRY_MAP.copy() + original_asgi = _function_registry.ASGI_FUNCTIONS.copy() + _function_registry.REGISTRY_MAP.clear() + _function_registry.ASGI_FUNCTIONS.clear() + yield + _function_registry.REGISTRY_MAP.clear() + _function_registry.REGISTRY_MAP.update(original_registry_map) + _function_registry.ASGI_FUNCTIONS.clear() + _function_registry.ASGI_FUNCTIONS.update(original_asgi) + + def test_get_function_signature(): test_cases = [ { diff --git a/tests/test_functions.py b/tests/test_functions.py index 534f4a88..bafb5e8b 100644 --- a/tests/test_functions.py +++ b/tests/test_functions.py @@ -495,7 +495,7 @@ def test_error_paths(http_trigger_client, path): def test_lazy_wsgi_app(monkeypatch, target, source, signature_type): actual_app_stub = pretend.stub() wsgi_app = pretend.call_recorder(lambda *a, **kw: actual_app_stub) - create_app = pretend.call_recorder(lambda *a: wsgi_app) + create_app = pretend.call_recorder(lambda *a, **kw: wsgi_app) monkeypatch.setattr(functions_framework, "create_app", create_app) # Test that it's lazy