Add warnings to FGA check / query response

Author: atainter

tests/test_fga.py

@@ -101,6 +101,89 @@ 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)
+
+    @pytest.fixture
+    def mock_check_warning_response(self):
+        return {
+            "result": "authorized",
+            "is_implicit": True,
+            "warnings": [
+                {
+                    "code": "missing_context_keys",
+                    "message": "Missing context keys",
+                    "keys": ["key1", "key2"],
+                }
+            ],
+        }
+
+    @pytest.fixture
+    def mock_query_warning_response(self):
+        return {
+            "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"],
+                }
+            ],
+        }
+
+    def test_check_with_warning(
+        self, mock_check_warning_response, mock_http_client_with_response
+    ):
+        mock_http_client_with_response(
+            self.http_client, mock_check_warning_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_check_warning_response
+
+    def test_query_with_warning(
+        self, mock_query_warning_response, mock_http_client_with_response
+    ):
+        mock_http_client_with_response(
+            self.http_client, mock_query_warning_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_query_warning_response
+
+
 class TestFGA:
     @pytest.fixture(autouse=True)
     def setup(self, sync_http_client_for_test):

workos/fga.py

@@ -10,7 +10,7 @@
     WarrantWrite,
     WarrantWriteOperation,
     WriteWarrantResponse,
-    WarrantQueryResult,
+    WarrantQueryResult, FGAWarning,
 )
 from workos.types.fga.list_filters import (
     AuthorizationResourceListFilters,
@@ -45,9 +45,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 +643,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/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,21 @@
+from typing import Sequence, Union, Literal, Annotated
+
+from pydantic import Field
+
+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]
+
+
+FGAWarning = Annotated[
+    Union[MissingContextKeysWarning, FGABaseWarning],
+    Field(discriminator='type')
+]