lamindb_setup 1.19.0__py3-none-any.whl → 1.19.1__py3-none-any.whl

lamindb_setup/core/_hub_client.py

@@ -1,288 +1,288 @@
-from __future__ import annotations
-
-import json
-import os
-from contextlib import contextmanager
-from datetime import datetime
-from typing import Literal
-from urllib.request import urlretrieve
-
-import httpx
-from httpx_retries import Retry, RetryTransport
-from lamin_utils import logger
-from supabase import Client, ClientOptions, create_client
-
-from ._settings_save import save_user_settings
-from ._settings_store import Connector
-
-
-def load_fallback_connector() -> Connector:
-    url = "https://lamin-site-assets.s3.amazonaws.com/connector.env"
-    connector_file, _ = urlretrieve(url)
-    return Connector.from_env_file(connector_file, "")
-
-
-PROD_URL = "https://hub.lamin.ai"
-PROD_ANON_KEY = "sb_publishable_YVa4h8hQ-yBhXpfa2cP39w_PhoLW6Nu"
-
-
-class Environment:
-    def __init__(self, fallback: bool = False):
-        lamin_env = os.getenv("LAMIN_ENV")
-        if lamin_env is None:
-            lamin_env = "prod"
-        # set public key
-        if lamin_env == "prod":
-            if not fallback:
-                url = PROD_URL
-                key = PROD_ANON_KEY
-            else:
-                connector = load_fallback_connector()
-                url = connector.url
-                key = connector.key
-        elif lamin_env == "staging":
-            url = "https://amvrvdwndlqdzgedrqdv.supabase.co"
-            key = "sb_publishable_amVjtilv_Yj4VmGLmxtq6A_sYlLoQx5"
-        elif lamin_env == "staging-test":
-            url = "https://iugyyajllqftbpidapak.supabase.co"
-            key = "sb_publishable_XmXroXqTLQw-eeT5kysCww_k8vJv-4L"
-        elif lamin_env == "prod-test":
-            url = "https://xtdacpwiqwpbxsatoyrv.supabase.co"
-            key = "sb_publishable_G-pyO5aW6VFErTzJyVvM5w_NAv1_Mf7"
-        else:
-            url = os.environ["SUPABASE_API_URL"]
-            key = os.environ["SUPABASE_ANON_KEY"]
-        self.lamin_env: str = lamin_env
-        self.supabase_api_url: str = url
-        self.supabase_anon_key: str = key
-
-
-DEFAULT_TIMEOUT = 12
-
-
-# needed to log retries
-class LogRetry(Retry):
-    def increment(self):
-        new = super().increment()
-        now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-        # new.attempts_made is the 1-based retry count
-        logger.warning(f"{now} HTTP retry attempt {new.attempts_made}/{new.total}")
-        return new
-
-
-# runs ~0.5s
-def connect_hub(
-    fallback_env: bool = False, client_options: ClientOptions | None = None
-) -> Client:
-    env = Environment(fallback=fallback_env)
-    if client_options is None:
-        client_options = ClientOptions(
-            auto_refresh_token=False,
-            function_client_timeout=DEFAULT_TIMEOUT,
-            postgrest_client_timeout=DEFAULT_TIMEOUT,
-        )
-    client = create_client(env.supabase_api_url, env.supabase_anon_key, client_options)
-    # needed to enable retries for http requests in supabase
-    # these are separate clients and need separate transports
-    transports = []
-    for _ in range(2):
-        transports.append(
-            RetryTransport(
-                retry=LogRetry(total=2, backoff_factor=0.2),
-                transport=httpx.HTTPTransport(verify=True, http2=True, trust_env=True),
-            )
-        )
-    # this overwrites transports of existing httpx clients
-    # if proxies are set, the default transports that were created on clients init
-    # will be used, irrespective of these re-settings
-    client.auth._http_client._transport = transports[0]
-    client.postgrest.session._transport = transports[1]
-    # POST is not retryable by default, but for our functions it should be safe to retry
-    client.functions._client._transport = RetryTransport(
-        retry=LogRetry(
-            total=2,
-            backoff_factor=0.2,
-            allowed_methods=[
-                "HEAD",
-                "GET",
-                "PUT",
-                "DELETE",
-                "OPTIONS",
-                "TRACE",
-                "POST",
-            ],
-        ),
-        transport=httpx.HTTPTransport(verify=True, http2=True, trust_env=True),
-    )
-    return client
-
-
-def connect_hub_with_auth(
-    fallback_env: bool = False,
-    renew_token: bool = False,
-    access_token: str | None = None,
-) -> Client:
-    hub = connect_hub(fallback_env=fallback_env)
-    if access_token is None:
-        from lamindb_setup import settings
-
-        if renew_token:
-            settings.user.access_token = get_access_token(
-                settings.user.email, settings.user.password, settings.user.api_key
-            )
-        access_token = settings.user.access_token
-    hub.postgrest.auth(access_token)
-    hub.functions.set_auth(access_token)
-    return hub
-
-
-# runs ~0.5s
-def get_access_token(
-    email: str | None = None, password: str | None = None, api_key: str | None = None
-):
-    hub = connect_hub()
-    try:
-        if api_key is not None:
-            auth_response = hub.functions.invoke(
-                "get-jwt-v1",
-                invoke_options={"body": {"api_key": api_key}},
-            )
-            return json.loads(auth_response)["accessToken"]
-        auth_response = hub.auth.sign_in_with_password(
-            {
-                "email": email,
-                "password": password,
-            }
-        )
-        return auth_response.session.access_token
-    except Exception as e:
-        # we need to log the problem here because the exception is usually caught outside
-        # in call_with_fallback_auth
-        logger.warning(f"failed to update your lamindb access token: {e}")
-        raise e
-    finally:
-        hub.auth.sign_out(options={"scope": "local"})
-
-
-def call_with_fallback_auth(
-    callable,
-    **kwargs,
-):
-    access_token = kwargs.pop("access_token", None)
-
-    if access_token is not None:
-        client = None
-        try:
-            client = connect_hub_with_auth(access_token=access_token)
-            result = callable(**kwargs, client=client)
-        finally:
-            if client is not None:
-                client.auth.sign_out(options={"scope": "local"})
-
-        return result
-
-    for renew_token, fallback_env in [(False, False), (True, False), (False, True)]:
-        client = None
-        try:
-            client = connect_hub_with_auth(
-                renew_token=renew_token, fallback_env=fallback_env
-            )
-            result = callable(**kwargs, client=client)
-            # we update access_token here
-            # because at this point the call has been successfully resolved
-            if renew_token:
-                from lamindb_setup import settings
-
-                # here settings.user contains an updated access_token
-                save_user_settings(settings.user)
-            break
-        # we use Exception here as the ways in which the client fails upon 401
-        # are not consistent and keep changing
-        # because we ultimately raise the error, it's OK I'd say
-        except Exception as e:
-            if fallback_env:
-                raise e
-        finally:
-            if client is not None:
-                client.auth.sign_out(options={"scope": "local"})
-
-    return result
-
-
-def call_with_fallback(
-    callable,
-    **kwargs,
-):
-    for fallback_env in [False, True]:
-        client = None
-        try:
-            client = connect_hub(fallback_env=fallback_env)
-            result = callable(**kwargs, client=client)
-            break
-        except Exception as e:
-            if fallback_env:
-                raise e
-        finally:
-            if client is not None:
-                # in case there was sign in
-                client.auth.sign_out(options={"scope": "local"})
-    return result
-
-
-@contextmanager
-def httpx_client():
-    client = None
-    try:
-        # local is used in tests
-        if os.environ.get("LAMIN_ENV", "prod") == "local":
-            from fastapi.testclient import TestClient
-            from laminhub_rest.main import app
-
-            client = TestClient(app)
-        else:
-            transport = RetryTransport(
-                retry=LogRetry(total=2, backoff_factor=0.2),
-                transport=httpx.HTTPTransport(verify=True, http2=True, trust_env=True),
-            )
-            # first we create a client to build the proxy map from the env variables
-            # if proxies are set, the default transports will be used
-            # otherwise the RetryTransport object that we assign below
-            client = httpx.Client(trust_env=True)
-            client._transport = transport
-        yield client
-    finally:
-        if client is not None:
-            client.close()
-
-
-def request_with_auth(
-    url: str,
-    method: Literal["get", "post", "put", "delete", "head", "options"],
-    access_token: str,
-    renew_token: bool = True,
-    **kwargs,
-):
-    headers = kwargs.pop("headers", {})
-    headers["Authorization"] = f"Bearer {access_token}"
-    timeout = kwargs.pop("timeout", DEFAULT_TIMEOUT)
-
-    with httpx_client() as client:
-        make_request = getattr(client, method)
-        response = make_request(url, headers=headers, timeout=timeout, **kwargs)
-        status_code = response.status_code
-        # update access_token and try again if failed
-        if not (200 <= status_code < 300) and renew_token:
-            from lamindb_setup import settings
-
-            logger.debug(f"{method} {url} failed: {status_code} {response.text}")
-
-            access_token = get_access_token(
-                settings.user.email, settings.user.password, settings.user.api_key
-            )
-
-            settings.user.access_token = access_token
-            save_user_settings(settings.user)
-
-            headers["Authorization"] = f"Bearer {access_token}"
-            response = make_request(url, headers=headers, timeout=timeout, **kwargs)
-    return response
+from __future__ import annotations
+
+import json
+import os
+from contextlib import contextmanager
+from datetime import datetime
+from typing import Literal
+from urllib.request import urlretrieve
+
+import httpx
+from httpx_retries import Retry, RetryTransport
+from lamin_utils import logger
+from supabase import Client, ClientOptions, create_client
+
+from ._settings_save import save_user_settings
+from ._settings_store import Connector
+
+
+def load_fallback_connector() -> Connector:
+    url = "https://lamin-site-assets.s3.amazonaws.com/connector.env"
+    connector_file, _ = urlretrieve(url)
+    return Connector.from_env_file(connector_file, "")
+
+
+PROD_URL = "https://hub.lamin.ai"
+PROD_ANON_KEY = "sb_publishable_YVa4h8hQ-yBhXpfa2cP39w_PhoLW6Nu"
+
+
+class Environment:
+    def __init__(self, fallback: bool = False):
+        lamin_env = os.getenv("LAMIN_ENV")
+        if lamin_env is None:
+            lamin_env = "prod"
+        # set public key
+        if lamin_env == "prod":
+            if not fallback:
+                url = PROD_URL
+                key = PROD_ANON_KEY
+            else:
+                connector = load_fallback_connector()
+                url = connector.url
+                key = connector.key
+        elif lamin_env == "staging":
+            url = "https://amvrvdwndlqdzgedrqdv.supabase.co"
+            key = "sb_publishable_amVjtilv_Yj4VmGLmxtq6A_sYlLoQx5"
+        elif lamin_env == "staging-test":
+            url = "https://iugyyajllqftbpidapak.supabase.co"
+            key = "sb_publishable_XmXroXqTLQw-eeT5kysCww_k8vJv-4L"
+        elif lamin_env == "prod-test":
+            url = "https://xtdacpwiqwpbxsatoyrv.supabase.co"
+            key = "sb_publishable_G-pyO5aW6VFErTzJyVvM5w_NAv1_Mf7"
+        else:
+            url = os.environ["SUPABASE_API_URL"]
+            key = os.environ["SUPABASE_ANON_KEY"]
+        self.lamin_env: str = lamin_env
+        self.supabase_api_url: str = url
+        self.supabase_anon_key: str = key
+
+
+DEFAULT_TIMEOUT = 12
+
+
+# needed to log retries
+class LogRetry(Retry):
+    def increment(self):
+        new = super().increment()
+        now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        # new.attempts_made is the 1-based retry count
+        logger.warning(f"{now} HTTP retry attempt {new.attempts_made}/{new.total}")
+        return new
+
+
+# runs ~0.5s
+def connect_hub(
+    fallback_env: bool = False, client_options: ClientOptions | None = None
+) -> Client:
+    env = Environment(fallback=fallback_env)
+    if client_options is None:
+        client_options = ClientOptions(
+            auto_refresh_token=False,
+            function_client_timeout=DEFAULT_TIMEOUT,
+            postgrest_client_timeout=DEFAULT_TIMEOUT,
+        )
+    client = create_client(env.supabase_api_url, env.supabase_anon_key, client_options)
+    # needed to enable retries for http requests in supabase
+    # these are separate clients and need separate transports
+    transports = []
+    for _ in range(2):
+        transports.append(
+            RetryTransport(
+                retry=LogRetry(total=2, backoff_factor=0.2),
+                transport=httpx.HTTPTransport(verify=True, http2=True, trust_env=True),
+            )
+        )
+    # this overwrites transports of existing httpx clients
+    # if proxies are set, the default transports that were created on clients init
+    # will be used, irrespective of these re-settings
+    client.auth._http_client._transport = transports[0]
+    client.postgrest.session._transport = transports[1]
+    # POST is not retryable by default, but for our functions it should be safe to retry
+    client.functions._client._transport = RetryTransport(
+        retry=LogRetry(
+            total=2,
+            backoff_factor=0.2,
+            allowed_methods=[
+                "HEAD",
+                "GET",
+                "PUT",
+                "DELETE",
+                "OPTIONS",
+                "TRACE",
+                "POST",
+            ],
+        ),
+        transport=httpx.HTTPTransport(verify=True, http2=True, trust_env=True),
+    )
+    return client
+
+
+def connect_hub_with_auth(
+    fallback_env: bool = False,
+    renew_token: bool = False,
+    access_token: str | None = None,
+) -> Client:
+    hub = connect_hub(fallback_env=fallback_env)
+    if access_token is None:
+        from lamindb_setup import settings
+
+        if renew_token:
+            settings.user.access_token = get_access_token(
+                settings.user.email, settings.user.password, settings.user.api_key
+            )
+        access_token = settings.user.access_token
+    hub.postgrest.auth(access_token)
+    hub.functions.set_auth(access_token)
+    return hub
+
+
+# runs ~0.5s
+def get_access_token(
+    email: str | None = None, password: str | None = None, api_key: str | None = None
+):
+    hub = connect_hub()
+    try:
+        if api_key is not None:
+            auth_response = hub.functions.invoke(
+                "get-jwt-v1",
+                invoke_options={"body": {"api_key": api_key}},
+            )
+            return json.loads(auth_response)["accessToken"]
+        auth_response = hub.auth.sign_in_with_password(
+            {
+                "email": email,
+                "password": password,
+            }
+        )
+        return auth_response.session.access_token
+    except Exception as e:
+        # we need to log the problem here because the exception is usually caught outside
+        # in call_with_fallback_auth
+        logger.warning(f"failed to update your lamindb access token: {e}")
+        raise e
+    finally:
+        hub.auth.sign_out(options={"scope": "local"})
+
+
+def call_with_fallback_auth(
+    callable,
+    **kwargs,
+):
+    access_token = kwargs.pop("access_token", None)
+
+    if access_token is not None:
+        client = None
+        try:
+            client = connect_hub_with_auth(access_token=access_token)
+            result = callable(**kwargs, client=client)
+        finally:
+            if client is not None:
+                client.auth.sign_out(options={"scope": "local"})
+
+        return result
+
+    for renew_token, fallback_env in [(False, False), (True, False), (False, True)]:
+        client = None
+        try:
+            client = connect_hub_with_auth(
+                renew_token=renew_token, fallback_env=fallback_env
+            )
+            result = callable(**kwargs, client=client)
+            # we update access_token here
+            # because at this point the call has been successfully resolved
+            if renew_token:
+                from lamindb_setup import settings
+
+                # here settings.user contains an updated access_token
+                save_user_settings(settings.user)
+            break
+        # we use Exception here as the ways in which the client fails upon 401
+        # are not consistent and keep changing
+        # because we ultimately raise the error, it's OK I'd say
+        except Exception as e:
+            if fallback_env:
+                raise e
+        finally:
+            if client is not None:
+                client.auth.sign_out(options={"scope": "local"})
+
+    return result
+
+
+def call_with_fallback(
+    callable,
+    **kwargs,
+):
+    for fallback_env in [False, True]:
+        client = None
+        try:
+            client = connect_hub(fallback_env=fallback_env)
+            result = callable(**kwargs, client=client)
+            break
+        except Exception as e:
+            if fallback_env:
+                raise e
+        finally:
+            if client is not None:
+                # in case there was sign in
+                client.auth.sign_out(options={"scope": "local"})
+    return result
+
+
+@contextmanager
+def httpx_client():
+    client = None
+    try:
+        # local is used in tests
+        if os.environ.get("LAMIN_ENV", "prod") == "local":
+            from fastapi.testclient import TestClient
+            from laminhub_rest.main import app
+
+            client = TestClient(app)
+        else:
+            transport = RetryTransport(
+                retry=LogRetry(total=2, backoff_factor=0.2),
+                transport=httpx.HTTPTransport(verify=True, http2=True, trust_env=True),
+            )
+            # first we create a client to build the proxy map from the env variables
+            # if proxies are set, the default transports will be used
+            # otherwise the RetryTransport object that we assign below
+            client = httpx.Client(trust_env=True)
+            client._transport = transport
+        yield client
+    finally:
+        if client is not None:
+            client.close()
+
+
+def request_with_auth(
+    url: str,
+    method: Literal["get", "post", "put", "delete", "head", "options"],
+    access_token: str,
+    renew_token: bool = True,
+    **kwargs,
+):
+    headers = kwargs.pop("headers", {})
+    headers["Authorization"] = f"Bearer {access_token}"
+    timeout = kwargs.pop("timeout", DEFAULT_TIMEOUT)
+
+    with httpx_client() as client:
+        make_request = getattr(client, method)
+        response = make_request(url, headers=headers, timeout=timeout, **kwargs)
+        status_code = response.status_code
+        # update access_token and try again if failed
+        if not (200 <= status_code < 300) and renew_token:
+            from lamindb_setup import settings
+
+            logger.debug(f"{method} {url} failed: {status_code} {response.text}")
+
+            access_token = get_access_token(
+                settings.user.email, settings.user.password, settings.user.api_key
+            )
+
+            settings.user.access_token = access_token
+            save_user_settings(settings.user)
+
+            headers["Authorization"] = f"Bearer {access_token}"
+            response = make_request(url, headers=headers, timeout=timeout, **kwargs)
+    return response