Add warnings field to FGA check / query response (#442)

* Add warnings to FGA check / query response

* Fix typing issues with Pydantic

* Fix python 3.8 issue

Author: atainter

tests/test_fga.py

@@ -101,6 +101,103 @@ def test_get_resource_401(self, mock_http_client_with_response):
             self.fga.get_resource(resource_type="test", resource_id="test")
 
 
+class TestWarnings:
+    @pytest.fixture(autouse=True)
+    def setup(self, sync_http_client_for_test):
+        self.http_client = sync_http_client_for_test
+        self.fga = FGA(http_client=self.http_client)
+
+    def test_check_with_warning(self, mock_http_client_with_response):
+        mock_response = {
+            "result": "authorized",
+            "is_implicit": True,
+            "warnings": [
+                {
+                    "code": "missing_context_keys",
+                    "message": "Missing context keys",
+                    "keys": ["key1", "key2"],
+                }
+            ],
+        }
+        mock_http_client_with_response(self.http_client, mock_response, 200)
+
+        response = self.fga.check(
+            op="any_of",
+            checks=[
+                WarrantCheckInput(
+                    resource_type="schedule",
+                    resource_id="schedule-A1",
+                    relation="viewer",
+                    subject=SubjectInput(resource_type="user", resource_id="user-A"),
+                )
+            ],
+        )
+        assert response.dict(exclude_none=True) == mock_response
+
+    def test_query_with_warning(self, mock_http_client_with_response):
+        mock_response = {
+            "object": "list",
+            "data": [
+                {
+                    "resource_type": "user",
+                    "resource_id": "richard",
+                    "relation": "member",
+                    "warrant": {
+                        "resource_type": "role",
+                        "resource_id": "developer",
+                        "relation": "member",
+                        "subject": {"resource_type": "user", "resource_id": "richard"},
+                    },
+                    "is_implicit": True,
+                }
+            ],
+            "list_metadata": {},
+            "warnings": [
+                {
+                    "code": "missing_context_keys",
+                    "message": "Missing context keys",
+                    "keys": ["key1", "key2"],
+                }
+            ],
+        }
+
+        mock_http_client_with_response(self.http_client, mock_response, 200)
+
+        response = self.fga.query(
+            q="select member of type user for permission:view-docs",
+            order="asc",
+            warrant_token="warrant_token",
+        )
+        assert response.dict(exclude_none=True) == mock_response
+
+    def test_check_with_generic_warning(self, mock_http_client_with_response):
+        mock_response = {
+            "result": "authorized",
+            "is_implicit": True,
+            "warnings": [
+                {
+                    "code": "generic",
+                    "message": "Generic warning",
+                }
+            ],
+        }
+
+        mock_http_client_with_response(self.http_client, mock_response, 200)
+
+        response = self.fga.check(
+            op="any_of",
+            checks=[
+                WarrantCheckInput(
+                    resource_type="schedule",
+                    resource_id="schedule-A1",
+                    relation="viewer",
+                    subject=SubjectInput(resource_type="user", resource_id="user-A"),
+                )
+            ],
+        )
+        assert response.dict(exclude_none=True) == mock_response
+
+
 class TestFGA:
     @pytest.fixture(autouse=True)
     def setup(self, sync_http_client_for_test):

workos/fga.py

@@ -11,6 +11,7 @@
     WarrantWriteOperation,
     WriteWarrantResponse,
     WarrantQueryResult,
+    FGAWarning,
 )
 from workos.types.fga.list_filters import (
     AuthorizationResourceListFilters,
@@ -45,9 +46,11 @@
 
 WarrantListResource = WorkOSListResource[Warrant, WarrantListFilters, ListMetadata]
 
-WarrantQueryListResource = WorkOSListResource[
-    WarrantQueryResult, WarrantQueryListFilters, ListMetadata
-]
+
+class WarrantQueryListResource(
+    WorkOSListResource[WarrantQueryResult, WarrantQueryListFilters, ListMetadata]
+):
+    warnings: Optional[Sequence[FGAWarning]] = None
 
 
 class FGAModule(Protocol):
@@ -641,5 +644,6 @@ def query(
         return WarrantQueryListResource(
             list_method=self.query,
             list_args=list_params,
+            warnings=response.get("warnings"),
             **ListPage[WarrantQueryResult](**response).model_dump(),
         )

workos/types/fga/__init__.py

@@ -2,3 +2,4 @@
 from .authorization_resource_types import *
 from .authorization_resources import *
 from .warrant import *
+from .warnings import *

workos/types/fga/check.py

@@ -3,6 +3,7 @@
 from workos.types.workos_model import WorkOSModel
 from workos.typing.literals import LiteralOrUntyped
 
+from .warnings import FGAWarning
 from .warrant import Subject, SubjectInput
 
 CheckOperation = Literal["any_of", "all_of", "batch"]
@@ -44,6 +45,7 @@ class CheckResponse(WorkOSModel):
     result: LiteralOrUntyped[CheckResult]
     is_implicit: bool
     debug_info: Optional[DebugInfo] = None
+    warnings: Optional[Sequence[FGAWarning]] = None
 
     def authorized(self) -> bool:
         return self.result == "authorized"

workos/types/fga/warnings.py

@@ -0,0 +1,33 @@
+from typing import Sequence, Union, Any, Dict, Literal
+from typing_extensions import Annotated
+
+from pydantic import BeforeValidator
+from pydantic_core.core_schema import ValidationInfo
+
+from workos.types.workos_model import WorkOSModel
+
+
+class FGABaseWarning(WorkOSModel):
+    code: str
+    message: str
+
+
+class MissingContextKeysWarning(FGABaseWarning):
+    code: Literal["missing_context_keys"]
+    keys: Sequence[str]
+
+
+def fga_warning_dispatch_validator(
+    value: Dict[str, Any], info: ValidationInfo
+) -> FGABaseWarning:
+    if value.get("code") == "missing_context_keys":
+        return MissingContextKeysWarning.model_validate(value)
+
+    # Fallback to the base warning model
+    return FGABaseWarning.model_validate(value)
+
+
+FGAWarning = Annotated[
+    Union[MissingContextKeysWarning, FGABaseWarning],
+    BeforeValidator(fga_warning_dispatch_validator),
+]