Merge pull request #92 from permitio/omer/per-10079-python-sdk-local-…

…facts-uploader Omer/per 10079 python sdk local facts uploader
permitio · Jun 19, 2024 · fde89ec · fde89ec
2 parents a46ce86 + 0e9230a
commit fde89ec
Show file tree

Hide file tree

Showing 9 changed files with 438 additions and 105 deletions.
diff --git a/permit/api/base.py b/permit/api/base.py
@@ -1,9 +1,11 @@
 import functools
-from typing import Callable, Optional, Type, TypeVar, Union
+from contextlib import contextmanager
+from typing import Callable, Iterable, Optional, Type, TypeVar, Union
 
 import aiohttp
 from aiohttp import ClientTimeout
 from loguru import logger
+from typing_extensions import Self
 
 from ..utils.pydantic_version import PYDANTIC_VERSION
 
@@ -12,7 +14,6 @@
 else:
     from pydantic.v1 import BaseModel, Extra, Field, parse_obj_as  # type: ignore
 
-
 from ..config import PermitConfig
 from ..exceptions import PermitContextError, handle_api_error, handle_client_error
 from .context import API_ACCESS_LEVELS, ApiContextLevel, ApiKeyAccessLevel
@@ -213,12 +214,18 @@ def __init__(self, config: PermitConfig):
         self.config = config
         self.__api_keys = self._build_http_client("/v2/api-key")
 
-    def _build_http_client(self, endpoint_url: str = "", **kwargs):
+    def _build_http_client(
+        self, endpoint_url: str = "", *, use_pdp: bool = False, **kwargs
+    ):
+        optional_headers = {}
+        if self.config.proxy_facts_via_pdp and self.config.facts_sync_timeout:
+            optional_headers["X-Wait-Timeout"] = str(self.config.facts_sync_timeout)
         client_config = ClientConfig(
-            base_url=f"{self.config.api_url}",
+            base_url=self.config.pdp if use_pdp else self.config.api_url,
             headers={
                 "Content-Type": "application/json",
                 "Authorization": f"bearer {self.config.token}",
+                **optional_headers,
             },
         )
         client_config = client_config.dict()

diff --git a/permit/api/relationship_tuples.py b/permit/api/relationship_tuples.py
@@ -29,12 +29,15 @@
 class RelationshipTuplesApi(BasePermitApi):
     @property
     def __relationship_tuples(self) -> SimpleHttpClient:
-        return self._build_http_client(
-            "/v2/facts/{proj_id}/{env_id}/relationship_tuples".format(
-                proj_id=self.config.api_context.project,
-                env_id=self.config.api_context.environment,
+        if self.config.proxy_facts_via_pdp:
+            return self._build_http_client("/facts/relationship_tuples", use_pdp=True)
+        else:
+            return self._build_http_client(
+                "/v2/facts/{proj_id}/{env_id}/relationship_tuples".format(
+                    proj_id=self.config.api_context.project,
+                    env_id=self.config.api_context.environment,
+                )
             )
-        )
 
     @required_permissions(ApiKeyAccessLevel.ENVIRONMENT_LEVEL_API_KEY)
     @required_context(ApiContextLevel.ENVIRONMENT)

diff --git a/permit/api/resource_instances.py b/permit/api/resource_instances.py
@@ -29,21 +29,29 @@
 class ResourceInstancesApi(BasePermitApi):
     @property
     def __resource_instances(self) -> SimpleHttpClient:
-        return self._build_http_client(
-            "/v2/facts/{proj_id}/{env_id}/resource_instances".format(
-                proj_id=self.config.api_context.project,
-                env_id=self.config.api_context.environment,
+        if self.config.proxy_facts_via_pdp:
+            return self._build_http_client("/facts/resource_instances", use_pdp=True)
+        else:
+            return self._build_http_client(
+                "/v2/facts/{proj_id}/{env_id}/resource_instances".format(
+                    proj_id=self.config.api_context.project,
+                    env_id=self.config.api_context.environment,
+                )
             )
-        )
 
     @property
     def __bulk_operations(self) -> SimpleHttpClient:
-        return self._build_http_client(
-            "/v2/facts/{proj_id}/{env_id}/bulk/resource_instances".format(
-                proj_id=self.config.api_context.project,
-                env_id=self.config.api_context.environment,
+        if self.config.proxy_facts_via_pdp:
+            return self._build_http_client(
+                "/facts/bulk/resource_instances", use_pdp=True
+            )
+        else:
+            return self._build_http_client(
+                "/v2/facts/{proj_id}/{env_id}/bulk/resource_instances".format(
+                    proj_id=self.config.api_context.project,
+                    env_id=self.config.api_context.environment,
+                )
             )
-        )
 
     @required_permissions(ApiKeyAccessLevel.ENVIRONMENT_LEVEL_API_KEY)
     @required_context(ApiContextLevel.ENVIRONMENT)

diff --git a/permit/api/role_assignments.py b/permit/api/role_assignments.py
@@ -28,12 +28,15 @@
 class RoleAssignmentsApi(BasePermitApi):
     @property
     def __role_assignments(self) -> SimpleHttpClient:
-        return self._build_http_client(
-            "/v2/facts/{proj_id}/{env_id}/role_assignments".format(
-                proj_id=self.config.api_context.project,
-                env_id=self.config.api_context.environment,
+        if self.config.proxy_facts_via_pdp:
+            return self._build_http_client("/facts/role_assignments", use_pdp=True)
+        else:
+            return self._build_http_client(
+                "/v2/facts/{proj_id}/{env_id}/role_assignments".format(
+                    proj_id=self.config.api_context.project,
+                    env_id=self.config.api_context.environment,
+                )
             )
-        )
 
     @required_permissions(ApiKeyAccessLevel.ENVIRONMENT_LEVEL_API_KEY)
     @required_context(ApiContextLevel.ENVIRONMENT)

diff --git a/permit/api/tenants.py b/permit/api/tenants.py
@@ -30,21 +30,27 @@
 class TenantsApi(BasePermitApi):
     @property
     def __tenants(self) -> SimpleHttpClient:
-        return self._build_http_client(
-            "/v2/facts/{proj_id}/{env_id}/tenants".format(
-                proj_id=self.config.api_context.project,
-                env_id=self.config.api_context.environment,
+        if self.config.proxy_facts_via_pdp:
+            return self._build_http_client("/facts/tenants", use_pdp=True)
+        else:
+            return self._build_http_client(
+                "/v2/facts/{proj_id}/{env_id}/tenants".format(
+                    proj_id=self.config.api_context.project,
+                    env_id=self.config.api_context.environment,
+                )
             )
-        )
 
     @property
     def __bulk_operations(self) -> SimpleHttpClient:
-        return self._build_http_client(
-            "/v2/facts/{proj_id}/{env_id}/bulk/tenants".format(
-                proj_id=self.config.api_context.project,
-                env_id=self.config.api_context.environment,
+        if self.config.proxy_facts_via_pdp:
+            return self._build_http_client("/facts/users", use_pdp=True)
+        else:
+            return self._build_http_client(
+                "/v2/facts/{proj_id}/{env_id}/bulk/tenants".format(
+                    proj_id=self.config.api_context.project,
+                    env_id=self.config.api_context.environment,
+                )
             )
-        )
 
     @required_permissions(ApiKeyAccessLevel.ENVIRONMENT_LEVEL_API_KEY)
     @required_context(ApiContextLevel.ENVIRONMENT)

diff --git a/permit/api/users.py b/permit/api/users.py
@@ -35,30 +35,39 @@
 class UsersApi(BasePermitApi):
     @property
     def __users(self) -> SimpleHttpClient:
-        return self._build_http_client(
-            "/v2/facts/{proj_id}/{env_id}/users".format(
-                proj_id=self.config.api_context.project,
-                env_id=self.config.api_context.environment,
+        if self.config.proxy_facts_via_pdp:
+            return self._build_http_client("/facts/users", use_pdp=True)
+        else:
+            return self._build_http_client(
+                "/v2/facts/{proj_id}/{env_id}/users".format(
+                    proj_id=self.config.api_context.project,
+                    env_id=self.config.api_context.environment,
+                )
             )
-        )
 
     @property
     def __role_assignments(self) -> SimpleHttpClient:
-        return self._build_http_client(
-            "/v2/facts/{proj_id}/{env_id}/role_assignments".format(
-                proj_id=self.config.api_context.project,
-                env_id=self.config.api_context.environment,
+        if self.config.proxy_facts_via_pdp:
+            return self._build_http_client("/facts/role_assignments", use_pdp=True)
+        else:
+            return self._build_http_client(
+                "/v2/facts/{proj_id}/{env_id}/role_assignments".format(
+                    proj_id=self.config.api_context.project,
+                    env_id=self.config.api_context.environment,
+                )
             )
-        )
 
     @property
     def __bulk_operations(self) -> SimpleHttpClient:
-        return self._build_http_client(
-            "/v2/facts/{proj_id}/{env_id}/bulk/users".format(
-                proj_id=self.config.api_context.project,
-                env_id=self.config.api_context.environment,
+        if self.config.proxy_facts_via_pdp:
+            return self._build_http_client("/facts/bulk/users", use_pdp=True)
+        else:
+            return self._build_http_client(
+                "/v2/facts/{proj_id}/{env_id}/bulk/users".format(
+                    proj_id=self.config.api_context.project,
+                    env_id=self.config.api_context.environment,
+                )
             )
-        )
 
     @required_permissions(ApiKeyAccessLevel.ENVIRONMENT_LEVEL_API_KEY)
     @required_context(ApiContextLevel.ENVIRONMENT)

diff --git a/permit/config.py b/permit/config.py
@@ -1,4 +1,4 @@
-from typing import List
+from typing import Optional
 
 from .api.context import ApiContext
 from .utils.pydantic_version import PYDANTIC_VERSION
@@ -69,6 +69,15 @@ class PermitConfig(BaseModel):
         None,
         description="The timeout in seconds for requests to the PDP.",
     )
+    proxy_facts_via_pdp: bool = Field(
+        False,
+        description="Create facts via the PDP API instead of using the default Permit REST API.",
+    )
+    facts_sync_timeout: Optional[float] = Field(
+        None,
+        description="The amount of time in seconds to wait for facts to be available "
+        "in the PDP cache before returning the response.",
+    )
 
     class Config:
         arbitrary_types_allowed = True
diff --git a/permit/permit.py b/permit/permit.py
@@ -1,7 +1,10 @@
 import json
-from typing import Optional
+from contextlib import contextmanager
+from typing import Generator, Optional
 
 from loguru import logger
+from pydantic import NonNegativeFloat
+from typing_extensions import Self
 
 from .api.api_client import PermitApiClient
 from .api.elements import ElementsApi
@@ -48,6 +51,33 @@ def config(self):
         """
         return self._config.copy()
 
+    @contextmanager
+    def wait_for_sync(self, timeout: float = 10.0) -> Generator[Self, None, None]:
+        """
+        Context manager that returns a client that is configured
+        to wait for facts to be synced before proceeding.
+
+
+        Args:
+            timeout: The amount of time in seconds to wait for facts to be available in the PDP
+            cache before returning the response.
+
+        Yields:
+            Permit: A Permit instance that is configured to wait for facts to be synced.
+
+        See Also:
+            https://docs.permit.io/how-to/manage-data/local-facts-uploader
+        """
+        if not self._config.proxy_facts_via_pdp:
+            logger.warning(
+                "Tried to wait for synced facts but proxy_facts_via_pdp is disabled, ignoring..."
+            )
+            yield self
+            return
+        contextualized_config = self.config  # this copies the config
+        contextualized_config.facts_sync_timeout = timeout
+        yield self.__class__(contextualized_config)
+
     @property
     def api(self) -> PermitApiClient:
         """