Support multiple security schemes in OpenAPI Client

fastagency/api/openapi/client.py

@@ -48,7 +48,7 @@
         self.registered_funcs: list[Callable[..., Any]] = []
         self.globals: dict[str, Any] = {}
 
-        self.security: dict[str, BaseSecurity] = {}
+        self.security: dict[str, list[BaseSecurity]] = {}
         self.security_params: dict[Optional[str], BaseSecurityParameters] = {}
 
     @staticmethod
@@ -106,18 +106,23 @@
             if security is None:
                 raise ValueError(f"Security is not set for '{name}'")
 
-            if not security.accept(security_params):
+            for match_security in security:
+                if match_security.accept(security_params):
+                    break
+            else:
                 raise ValueError(
                     f"Security parameters {security_params} do not match security {security}"
                 )
 
         self.security_params[name] = security_params
 
-    def _get_security_params(self, name: str) -> Optional[BaseSecurityParameters]:
+    def _get_security_params(
+        self, name: str
+    ) -> tuple[Optional[BaseSecurityParameters], Optional[BaseSecurity]]:
         # check if security is set for the method
         security = self.security.get(name)
-        if security is None:
-            return None
+        if not security:
+            return None, None
 
         security_params = self.security_params.get(name)
         if security_params is None:
@@ -129,19 +134,22 @@
                 )
 
         # check if security matches security parameters
-        if not security.accept(security_params):
+        for match_security in security:
+            if match_security.accept(security_params):
+                break
+        else:
             raise ValueError(
                 f"Security parameters {security_params} do not match security {security}"
             )
 
-        return security_params
+        return security_params, match_security
 
     def _request(
         self,
         method: Literal["put", "get", "post", "delete"],
         path: str,
         description: Optional[str] = None,
-        security: Optional[BaseSecurity] = None,
+        security: Optional[list[BaseSecurity]] = None,
         **kwargs: Any,
     ) -> Callable[..., dict[str, Any]]:
         def decorator(func: Callable[..., Any]) -> Callable[..., dict[str, Any]]:
@@ -156,13 +164,13 @@
 
                 security = self.security.get(name)
                 if security is not None:
-                    security_params = self._get_security_params(name)
+                    security_params, matched_security = self._get_security_params(name)
                     if security_params is None:
                         raise ValueError(
                             f"Security parameters are not set for '{name}'"
                         )
                     else:
-                        security_params.apply(params, body_dict, security)
+                        security_params.apply(params, body_dict, matched_security)  # type: ignore [arg-type]
 
                 response = getattr(requests, method)(url, params=params, **body_dict)
                 return response.json()  # type: ignore [no-any-return]

templates/main.jinja2

@@ -38,11 +38,11 @@ app = OpenAPI(
     , tags={{operation.tags}}
     {% endif %}
     {% if operation.security %}
-    {% for security in operation.security %}
+    , security=[{% for security in operation.security %}
     {% for key, value in security.items() %}
-    , security={{security_parameters[key]}}
-    {% endfor %}
+    {{security_parameters[key]}},
     {% endfor %}
+    {% endfor %}]
     {% endif %}
     )
 def {{operation.function_name}}({{operation.snake_case_arguments}}

tests/api/openapi/security/conftest.py

@@ -5,28 +5,42 @@
 
 import pytest
 import uvicorn
-from fastapi import Depends, FastAPI, Query
+from fastapi import Depends, FastAPI, HTTPException, Query, status
 from fastapi.responses import JSONResponse
-from fastapi.security import APIKeyHeader
+from fastapi.security import APIKeyCookie, APIKeyHeader
 
 from ....conftest import Server, find_free_port
 
 
 def create_secure_fastapi_app(host: str, port: int) -> FastAPI:
     app = FastAPI(servers=[{"url": f"http://{host}:{port}"}])
 
-    header_scheme = APIKeyHeader(name="x-key")
-
-    universal_api_key = "super secret key"  # pragma: allowlist secret
+    api_key = "super secret key"  # pragma: allowlist secret
+    api_key_name = "access_token"  # pragma: allowlist secret
+
+    header_scheme = APIKeyHeader(name=api_key_name, auto_error=False)
+    cookie_scheme = APIKeyCookie(name=api_key_name, auto_error=False)
+
+    async def get_api_key(
+        api_key_header: str = Depends(header_scheme),
+        api_key_cookie: str = Depends(cookie_scheme),
+    ) -> str:
+        if api_key_header == api_key:
+            return api_key_header
+        elif api_key_cookie == api_key:
+            return api_key_cookie
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid or missing API Key",
+        )
 
     @app.get("/items/")
     async def read_items(
         city: Annotated[str, Query(description="city for which forecast is requested")],
-        key: str = Depends(header_scheme),
+        api_key: str = Depends(get_api_key),
     ) -> JSONResponse:
-        is_authenticated = key == universal_api_key
-        content = {"is_authenticated": is_authenticated}
-        status_code = 200 if is_authenticated else 403
+        content = {"api_key": api_key}
+        status_code = 200
         return JSONResponse(status_code=status_code, content=content)
 
     return app

tests/api/openapi/security/expected_main_gen.txt

@@ -8,7 +8,7 @@ from typing import *
 from typing import Any, Union
 
 from fastagency.api.openapi import OpenAPI
-from fastagency.api.openapi.security import APIKeyHeader
+from fastagency.api.openapi.security import APIKeyCookie, APIKeyHeader
 
 from models_gen import HTTPValidationError
 
@@ -23,7 +23,10 @@ app = OpenAPI(
     '/items/',
     response_model=Any,
     responses={'422': {'model': HTTPValidationError}},
-    security=APIKeyHeader(name="x-key"),
+    security=[
+        APIKeyHeader(name="access_token"),
+        APIKeyCookie(name="access_token"),
+    ],
 )
 def read_items_items__get(
     city: Annotated[str, """city for which forecast is requested"""]

tests/api/openapi/security/expected_openapi.json

@@ -17,6 +17,9 @@
                 "security": [
                     {
                         "APIKeyHeader": []
+                    },
+                    {
+                        "APIKeyCookie": []
                     }
                 ],
                 "parameters": [
@@ -108,7 +111,12 @@
             "APIKeyHeader": {
                 "type": "apiKey",
                 "in": "header",
-                "name": "x-key"
+                "name": "access_token"
+            },
+            "APIKeyCookie": {
+                "type": "apiKey",
+                "in": "cookie",
+                "name": "access_token"
             }
         }
     }

tests/api/openapi/security/test_security.py

@@ -6,7 +6,7 @@
 import requests
 
 from fastagency.api.openapi import OpenAPI
-from fastagency.api.openapi.security import APIKeyHeader
+from fastagency.api.openapi.security import APIKeyCookie, APIKeyHeader
 
 
 def test_secure_app_openapi(secure_fastapi_url: str) -> None:
@@ -56,6 +56,7 @@ def test_generate_client(secure_fastapi_url: str) -> None:
         with expected_models_gen_path.open() as f:
             expected_models_gen = f.readlines()[4:]
 
+        # print(actual_main_gen_txt)
         assert actual_main_gen_txt == expected_main_gen_txt
         assert actual_models_gen == expected_models_gen
 
@@ -81,15 +82,24 @@ def test_import_and_call_generate_client(secure_fastapi_url: str) -> None:
 
         assert generated_client_app.security != {}, generated_client_app.security
 
+        api_key = "super secret key"  # pragma: allowlist secret
+
         # set global security params for all methods
-        # generated_client_app.set_security_params(APIKeyHeader.Parameters(value="super secret key"))
+        # generated_client_app.set_security_params(APIKeyHeader.Parameters(value=api_key))
 
         # or set security params for a specific method
         generated_client_app.set_security_params(
-            APIKeyHeader.Parameters(value="super secret key"),
+            APIKeyHeader.Parameters(value=api_key),
             "read_items_items__get",
         )
-
         # no security params added to the signature of the method
         client_resp = read_items_items__get(city="New York")
-        assert client_resp == {"is_authenticated": True}
+        assert client_resp == {"api_key": api_key}
+
+        # Test with cookie security
+        generated_client_app.set_security_params(
+            APIKeyCookie.Parameters(value=api_key),
+            "read_items_items__get",
+        )
+        client_resp = read_items_items__get(city="New York")
+        assert client_resp == {"api_key": api_key}