Add query and batch check implementations

tests/test_fga.py

@@ -467,3 +467,79 @@ def test_check_with_debug_info(
             debug=True,
         )
         assert response.dict(exclude_none=True) == mock_check_response_with_debug_info
+
+    @pytest.fixture
+    def mock_batch_check_response(self):
+        return [
+            {"result": "authorized", "is_implicit": True},
+            {"result": "not_authorized", "is_implicit": True},
+        ]
+
+    def test_check_batch(
+        self, mock_batch_check_response, mock_http_client_with_response
+    ):
+        mock_http_client_with_response(self.http_client, mock_batch_check_response, 200)
+
+        response = self.fga.check_batch(
+            checks=[
+                WarrantCheck(
+                    resource_type="schedule",
+                    resource_id="schedule-A1",
+                    relation="viewer",
+                    subject=Subject(resource_type="user", resource_id="user-A"),
+                ),
+                WarrantCheck(
+                    resource_type="schedule",
+                    resource_id="schedule-A1",
+                    relation="editor",
+                    subject=Subject(resource_type="user", resource_id="user-B"),
+                ),
+            ]
+        )
+
+        assert [
+            r.dict(exclude_none=True) for r in response
+        ] == mock_batch_check_response
+
+    @pytest.fixture
+    def mock_query_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,
+                },
+                {
+                    "resource_type": "user",
+                    "resource_id": "tom",
+                    "relation": "member",
+                    "warrant": {
+                        "resource_type": "role",
+                        "resource_id": "manager",
+                        "relation": "member",
+                        "subject": {"resource_type": "user", "resource_id": "tom"},
+                    },
+                    "is_implicit": True,
+                },
+            ],
+            "list_metadata": {},
+        }
+
+    def test_query(self, mock_query_response, mock_http_client_with_response):
+        mock_http_client_with_response(self.http_client, mock_query_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_response

workos/fga.py

@@ -11,8 +11,13 @@
     WarrantWrite,
     WarrantWriteOperation,
     WriteWarrantResponse,
+    WarrantQueryResult,
+)
+from workos.types.fga.list_filters import (
+    ResourceListFilters,
+    WarrantListFilters,
+    QueryListFilters,
 )
-from workos.types.fga.list_filters import ResourceListFilters, WarrantListFilters
 from workos.types.list_resource import (
     ListArgs,
     ListMetadata,
@@ -38,6 +43,10 @@
 
 WarrantListResource = WorkOsListResource[Warrant, WarrantListFilters, ListMetadata]
 
+QueryListResource = WorkOsListResource[
+    WarrantQueryResult, QueryListFilters, ListMetadata
+]
+
 
 class FGAModule(Protocol):
     def get_resource(self, *, resource_type: str, resource_id: str) -> Resource: ...
@@ -122,6 +131,26 @@ def check(
         warrant_token: Optional[str] = None,
     ) -> CheckResponse: ...
 
+    def check_batch(
+        self,
+        *,
+        checks: List[WarrantCheck],
+        debug: bool = False,
+        warrant_token: Optional[str] = None,
+    ) -> List[CheckResponse]: ...
+
+    def query(
+        self,
+        *,
+        q: str,
+        limit: int = DEFAULT_RESPONSE_LIMIT,
+        order: PaginationOrder = "desc",
+        before: Optional[str] = None,
+        after: Optional[str] = None,
+        context: Optional[Dict[str, Any]] = None,
+        warrant_token: Optional[str] = None,
+    ) -> WorkOsListResource[WarrantQueryResult, QueryListFilters, ListMetadata]: ...
+
 
 class FGA(FGAModule):
     _http_client: SyncHTTPClient
@@ -401,3 +430,65 @@ def check(
         )
 
         return CheckResponse.model_validate(response)
+
+    def check_batch(
+        self,
+        *,
+        checks: List[WarrantCheck],
+        debug: bool = False,
+        warrant_token: Optional[str] = None,
+    ) -> List[CheckResponse]:
+        if not checks:
+            raise ValueError("Incomplete arguments: No checks provided")
+
+        body = {
+            "checks": [check.dict() for check in checks],
+            "debug": debug,
+        }
+
+        response = self._http_client.request(
+            "fga/v1/check",
+            method=REQUEST_METHOD_POST,
+            token=workos.api_key,
+            json=body,
+            headers={"Warrant-Token": warrant_token} if warrant_token else None,
+        )
+
+        return [CheckResponse.model_validate(check) for check in response]
+
+    def query(
+        self,
+        *,
+        q: str,
+        limit: int = DEFAULT_RESPONSE_LIMIT,
+        order: PaginationOrder = "desc",
+        before: Optional[str] = None,
+        after: Optional[str] = None,
+        context: Optional[Dict[str, Any]] = None,
+        warrant_token: Optional[str] = None,
+    ) -> WorkOsListResource[WarrantQueryResult, QueryListFilters, ListMetadata]:
+        list_params: QueryListFilters = {
+            "q": q,
+            "limit": limit,
+            "order": order,
+            "before": before,
+            "after": after,
+            "context": context,
+        }
+
+        response = self._http_client.request(
+            "fga/v1/query",
+            method=REQUEST_METHOD_GET,
+            token=workos.api_key,
+            params=list_params,
+            headers={"Warrant-Token": warrant_token} if warrant_token else None,
+        )
+
+        # A workaround to add warrant_token to the list_args for the ListResource iterator
+        list_params["warrant_token"] = warrant_token
+
+        return WorkOsListResource[WarrantQueryResult, QueryListFilters, ListMetadata](
+            list_method=self.query,
+            list_args=list_params,
+            **ListPage[WarrantQueryResult](**response).model_dump(),
+        )

workos/types/fga/list_filters.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Dict, Any
 
 from workos.types.list_resource import ListArgs
 
@@ -16,3 +16,9 @@ class WarrantListFilters(ListArgs, total=False):
     subject_id: Optional[str]
     subject_relation: Optional[str]
     warrant_token: Optional[str]
+
+
+class QueryListFilters(ListArgs, total=False):
+    q: Optional[str]
+    context: Optional[Dict[str, Any]]
+    warrant_token: Optional[str]

workos/types/fga/warrant.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Literal, Optional
+from typing import Literal, Optional, Dict, Any
 
 from workos.types.workos_model import WorkOSModel
 
@@ -37,3 +37,12 @@ class WarrantWrite(WorkOSModel):
     relation: str
     subject: Subject
     policy: Optional[str] = None
+
+
+class WarrantQueryResult(WorkOSModel):
+    resource_type: str
+    resource_id: str
+    relation: str
+    warrant: Warrant
+    is_implicit: bool
+    meta: Optional[Dict[str, Any]] = None

workos/types/list_resource.py

@@ -23,7 +23,7 @@
     DirectoryUserWithGroups,
 )
 from workos.types.events import Event
-from workos.types.fga import Warrant, Resource, ResourceType
+from workos.types.fga import Warrant, Resource, ResourceType, WarrantQueryResult
 from workos.types.mfa import AuthenticationFactor
 from workos.types.organizations import Organization
 from workos.types.sso import ConnectionWithDomains
@@ -46,6 +46,7 @@
     ResourceType,
     User,
     Warrant,
+    WarrantQueryResult,
 )