dex-python-sdk 1.15.0__tar.gz → 1.16.0__tar.gz

{dex_python_sdk-1.15.0 → dex_python_sdk-1.16.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dex-python-sdk
-Version: 1.15.0
+Version: 1.16.0
 Requires-Dist: grpcio>=1.83.0
 Requires-Dist: grpcio-status>=1.83.0
 Requires-Dist: protobuf>=7.35.1
@@ -25,6 +25,47 @@ checks its mtime and size at most once per second, and reloads it after rotation
 Whitespace around the token is removed. An unreadable or empty file fails calls
 until it is restored; the SDK does not send an old token.
 
+Use `ApiKeyCredential` when the application holds a long-lived API key and a
+token endpoint issues short-lived access tokens scoped to a server ID. Supply
+the token URL, server ID, and API key explicitly; the server ID is never derived
+from the endpoint host. The endpoint must use HTTPS, except loopback HTTP on
+`localhost`, `127.0.0.0/8`, or `::1` for local development.
+Loopback HTTP bypasses environment proxies so the API key stays on the host.
+HTTPS uses configured proxies with TLS through CONNECT. Token endpoint redirects
+are refused, including redirects that preserve the POST method.
+
+```python
+from dex import ApiKeyCredential, Client, ClientOptions
+
+credential = ApiKeyCredential(
+    "https://auth.example.com/token", "dex-server", "application-api-key"
+)
+client = Client(registry, cache, ClientOptions(credential=credential))
+try:
+    client.health_check()
+finally:
+    client.close()
+    credential.close()
+```
+
+The first call exchanges the key for a token. After 80% of its lifetime, a
+single background refresh begins while calls continue using the valid token.
+Callers waiting for an expired token share the exchange. Failed refreshes retry
+no more often than every five seconds while the previous token remains valid.
+Each exchange has a 10-second deadline by default, configurable through
+`exchange_timeout`. It covers connecting, sending, reading, and parsing. A Dex
+`UNAUTHENTICATED` response discards the rejected token
+and retries a unary call once with a fresh token. Pull Worker streams reconnect
+once immediately with a fresh token; a second rejection uses the normal backoff.
+The credential is thread-safe.
+It redacts the API key, issued access tokens, and JWT-shaped strings from
+exchange errors, including endpoint details and transport causes. Client
+calls raise `ApiKeyExchangeError` unchanged; pull Workers log it and retry.
+Closing the credential is optional. After all Clients and Workers sharing it
+are closed, `close()` prevents further token calls and waits at most until the
+current exchange deadline. An uninterruptible network operation may continue
+in its daemon thread after that deadline.
+
 ```python
 from dex import Client, ClientOptions, FileCredential, Worker, WorkerOptions
 

{dex_python_sdk-1.15.0 → dex_python_sdk-1.16.0}/README.md

@@ -9,6 +9,47 @@ checks its mtime and size at most once per second, and reloads it after rotation
 Whitespace around the token is removed. An unreadable or empty file fails calls
 until it is restored; the SDK does not send an old token.
 
+Use `ApiKeyCredential` when the application holds a long-lived API key and a
+token endpoint issues short-lived access tokens scoped to a server ID. Supply
+the token URL, server ID, and API key explicitly; the server ID is never derived
+from the endpoint host. The endpoint must use HTTPS, except loopback HTTP on
+`localhost`, `127.0.0.0/8`, or `::1` for local development.
+Loopback HTTP bypasses environment proxies so the API key stays on the host.
+HTTPS uses configured proxies with TLS through CONNECT. Token endpoint redirects
+are refused, including redirects that preserve the POST method.
+
+```python
+from dex import ApiKeyCredential, Client, ClientOptions
+
+credential = ApiKeyCredential(
+    "https://auth.example.com/token", "dex-server", "application-api-key"
+)
+client = Client(registry, cache, ClientOptions(credential=credential))
+try:
+    client.health_check()
+finally:
+    client.close()
+    credential.close()
+```
+
+The first call exchanges the key for a token. After 80% of its lifetime, a
+single background refresh begins while calls continue using the valid token.
+Callers waiting for an expired token share the exchange. Failed refreshes retry
+no more often than every five seconds while the previous token remains valid.
+Each exchange has a 10-second deadline by default, configurable through
+`exchange_timeout`. It covers connecting, sending, reading, and parsing. A Dex
+`UNAUTHENTICATED` response discards the rejected token
+and retries a unary call once with a fresh token. Pull Worker streams reconnect
+once immediately with a fresh token; a second rejection uses the normal backoff.
+The credential is thread-safe.
+It redacts the API key, issued access tokens, and JWT-shaped strings from
+exchange errors, including endpoint details and transport causes. Client
+calls raise `ApiKeyExchangeError` unchanged; pull Workers log it and retry.
+Closing the credential is optional. After all Clients and Workers sharing it
+are closed, `close()` prevents further token calls and waits at most until the
+current exchange deadline. An uninterruptible network operation may continue
+in its daemon thread after that deadline.
+
 ```python
 from dex import Client, ClientOptions, FileCredential, Worker, WorkerOptions
 

{dex_python_sdk-1.15.0 → dex_python_sdk-1.16.0}/dex/__init__.py

@@ -24,7 +24,6 @@ from dex.blob_cache import BlobCache, BlobCacheConfig, open_blob_cache
 from dex.channel import Channel, ChannelMap, ChannelMapLoad, ChannelMessage
 from dex.client import Client
 from dex.client_options import ClientOptions
-from dex.credential import Credential, FileCredential
 from dex.codec import (
     BOOL,
     BYTES,
@@ -39,6 +38,12 @@ from dex.codec import (
 )
 from dex.condition import ConditionCombination
 from dex.context import AsyncContext, Context, RecoveryErrorInfo
+from dex.credential import (
+    ApiKeyCredential,
+    ApiKeyExchangeError,
+    Credential,
+    FileCredential,
+)
 from dex.flow import Flow, PersistenceSchema, Registry, RPCResult, rpc
 from dex.flow_config import ActiveStepSearchMode, FlowConfig
 from dex.flow_info import (
@@ -62,6 +67,12 @@ from dex.flow_options import (
     TimeTravelType,
 )
 from dex.flow_result import FlowResult, StepCompletion
+from dex.pull_registration import (
+    PullQueueRegistrationState,
+    PullRegistrationError,
+    PullRegistrationState,
+)
+from dex.rpc_invoke_options import RPCInvokeOptions
 from dex.runtime_errors import (
     AttributeMapNotLoadedError,
     ChannelMessageNotFoundError,
@@ -80,7 +91,6 @@ from dex.runtime_errors import (
     ValueMappingError,
     WorkerInvocationError,
 )
-from dex.rpc_invoke_options import RPCInvokeOptions
 from dex.step import (
     RetryPolicy,
     Step,
@@ -113,11 +123,6 @@ from dex.timer import Timer
 from dex.wait import Wait
 from dex.wait_options import WaitForAttributeOptions, WaitForStepCompletionOptions
 from dex.worker import Worker
-from dex.pull_registration import (
-    PullRegistrationError,
-    PullRegistrationState,
-    PullQueueRegistrationState,
-)
 from dex.worker_options import PullWorkerOptions, WorkerOptions, WorkerTarget
 
 __all__ = [
@@ -138,6 +143,8 @@ __all__ = [
     "AsyncClient",
     "AsyncBufferedTextStream",
     "AsyncWorker",
+    "ApiKeyCredential",
+    "ApiKeyExchangeError",
     "BlobCache",
     "BlobCacheConfig",
     "BufferedTextStream",

{dex_python_sdk-1.15.0 → dex_python_sdk-1.16.0}/dex/_credential_transport.py

@@ -16,6 +16,7 @@ import asyncio
 import inspect
 import queue
 import threading
+import time
 from collections.abc import Callable
 from concurrent.futures import Future
 from typing import Any, NamedTuple, cast
@@ -168,6 +169,38 @@ def credential_deadline_error() -> grpc.aio.AioRpcError:
     )
 
 
+class _SyncCredentialDeadlineError(grpc.RpcError):
+    def code(self) -> grpc.StatusCode:
+        return grpc.StatusCode.DEADLINE_EXCEEDED
+
+    def details(self) -> str:
+        return "credential lookup exceeded RPC deadline"
+
+    def trailing_metadata(self) -> tuple[Any, ...]:
+        return ()
+
+
+def _bounded_credential_metadata(
+    credential: Credential, metadata: Any, timeout: float | None
+) -> Any:
+    if timeout is not None and timeout <= 0:
+        raise _SyncCredentialDeadlineError()
+    bounded_token = getattr(credential, "_token_with_timeout", None)
+    if not callable(bounded_token) or timeout is None:
+        return credential_metadata(credential, metadata)
+    try:
+        token = bounded_token(timeout)
+    except TimeoutError:
+        raise _SyncCredentialDeadlineError() from None
+    return tuple(metadata or ()) + (("authorization", "Bearer " + token),)
+
+
+def _reject_credential_token(credential: Credential, metadata: Any) -> None:
+    reject_token = getattr(credential, "_reject_token", None)
+    if callable(reject_token):
+        reject_token(metadata[-1][1][7:])
+
+
 class _SyncCredentialInterceptor(grpc.UnaryUnaryClientInterceptor):
     def __init__(self, credential: Credential) -> None:
         self._credential = credential
@@ -178,15 +211,49 @@ class _SyncCredentialInterceptor(grpc.UnaryUnaryClientInterceptor):
         client_call_details: grpc.ClientCallDetails,
         request: Any,
     ) -> Any:
+        started_at = time.monotonic()
+        timeout = client_call_details.timeout
+        metadata = _bounded_credential_metadata(
+            self._credential, client_call_details.metadata, timeout
+        )
+        if timeout is not None:
+            timeout -= time.monotonic() - started_at
+            if timeout <= 0:
+                raise _SyncCredentialDeadlineError()
         details = _SyncCallDetails(
             client_call_details.method,
-            client_call_details.timeout,
-            credential_metadata(self._credential, client_call_details.metadata),
+            timeout,
+            metadata,
             client_call_details.credentials,
             client_call_details.wait_for_ready,
             client_call_details.compression,
         )
-        return continuation(details, request)
+        try:
+            response = continuation(details, request)
+        except grpc.RpcError as error:
+            if error.code() != grpc.StatusCode.UNAUTHENTICATED:
+                raise
+            if not callable(getattr(self._credential, "_reject_token", None)):
+                raise
+        else:
+            if response.code() != grpc.StatusCode.UNAUTHENTICATED:
+                return response
+            if not callable(getattr(self._credential, "_reject_token", None)):
+                return response
+        _reject_credential_token(self._credential, details.metadata)
+        if client_call_details.timeout is not None:
+            timeout = client_call_details.timeout - (time.monotonic() - started_at)
+        retry_details = details._replace(
+            metadata=_bounded_credential_metadata(
+                self._credential, client_call_details.metadata, timeout
+            )
+        )
+        if client_call_details.timeout is not None:
+            timeout = client_call_details.timeout - (time.monotonic() - started_at)
+            if timeout <= 0:
+                raise _SyncCredentialDeadlineError()
+            retry_details = retry_details._replace(timeout=timeout)
+        return continuation(retry_details, request)
 
 
 class _AsyncCredentialInterceptor(grpc.aio.UnaryUnaryClientInterceptor):
@@ -216,7 +283,39 @@ class _AsyncCredentialInterceptor(grpc.aio.UnaryUnaryClientInterceptor):
             client_call_details.credentials,
             client_call_details.wait_for_ready,
         )
-        return await continuation(details, request)
+        call = await continuation(details, request)
+        try:
+            await call
+        except grpc.aio.AioRpcError as error:
+            if error.code() != grpc.StatusCode.UNAUTHENTICATED:
+                raise
+            if not callable(getattr(self._runner._credential, "_reject_token", None)):
+                raise
+            _reject_credential_token(self._runner._credential, metadata)
+            if client_call_details.timeout is not None:
+                remaining_timeout = client_call_details.timeout - (
+                    loop.time() - started_at
+                )
+                if remaining_timeout <= 0:
+                    raise credential_deadline_error() from None
+            retry_metadata = await async_credential_metadata(
+                self._runner, client_call_details.metadata, remaining_timeout
+            )
+            if client_call_details.timeout is not None:
+                remaining_timeout = client_call_details.timeout - (
+                    loop.time() - started_at
+                )
+                if remaining_timeout <= 0:
+                    raise credential_deadline_error() from None
+            retry_details = grpc.aio.ClientCallDetails(
+                client_call_details.method,
+                remaining_timeout,
+                retry_metadata,
+                client_call_details.credentials,
+                client_call_details.wait_for_ready,
+            )
+            return await continuation(retry_details, request)
+        return call
 
 
 def sync_channel(server_address: str, credential: Credential | None) -> grpc.Channel:

{dex_python_sdk-1.15.0 → dex_python_sdk-1.16.0}/dex/_pull_worker.py

@@ -22,7 +22,11 @@ from typing import Any
 
 import grpc
 
-from dex._credential_transport import AsyncCredentialRunner, async_credential_metadata
+from dex._credential_transport import (
+    AsyncCredentialRunner,
+    _reject_credential_token,
+    async_credential_metadata,
+)
 from dex._grpc_errors import _worker_error_status
 from dex._server_protocol import sdk_version
 from dex.dexpb import dex_pb2 as pb
@@ -314,6 +318,7 @@ class PullStream:
         self._call: Any = None
         self._send_lock = asyncio.Lock()
         self._is_reconnecting = False
+        self._stream_metadata: Any = None
         self._deadline_handle: asyncio.TimerHandle | None = None
         self._supports_unknown_name_rejection = False
 
@@ -322,6 +327,7 @@ class PullStream:
 
     async def run(self) -> None:
         backoff = 0.1
+        has_retried_unauthenticated = False
         while not self.runtime.is_stopping:
             self.is_draining = False
             was_registered = False
@@ -336,12 +342,35 @@ class PullStream:
                 )
             finally:
                 code = error.code() if error is not None else None
+                if self._is_registered:
+                    has_retried_unauthenticated = False
+                credential = (
+                    self.runtime._credential_runner._credential
+                    if self.runtime._credential_runner is not None
+                    else None
+                )
+                can_reject_token = (
+                    code is grpc.StatusCode.UNAUTHENTICATED
+                    and self._stream_metadata is not None
+                    and credential is not None
+                    and callable(getattr(credential, "_reject_token", None))
+                )
+                should_retry_unauthenticated = (
+                    can_reject_token and not has_retried_unauthenticated
+                )
+                if can_reject_token and credential is not None:
+                    _reject_credential_token(credential, self._stream_metadata)
                 rejection = self.runtime.registration.record_stream_result(
                     self.queue,
                     self._is_registered,
-                    error if code in _PERMANENT_CODES else None,
+                    (
+                        error
+                        if code in _PERMANENT_CODES and not should_retry_unauthenticated
+                        else None
+                    ),
                 )
                 self._is_registered = False
+                self._stream_metadata = None
                 if (
                     rejection is not None
                     and self.runtime._stop_on_permanent_stream_rejection
@@ -357,6 +386,11 @@ class PullStream:
                 or self._is_retiring
             ):
                 break
+            if should_retry_unauthenticated:
+                has_retried_unauthenticated = True
+                continue
+            if code is grpc.StatusCode.UNAUTHENTICATED:
+                has_retried_unauthenticated = False
             if self._is_reconnecting:
                 self._is_reconnecting = False
                 backoff = 0.1
@@ -390,6 +424,7 @@ class PullStream:
         if self.runtime.is_stopping:
             return False
         self._supports_unknown_name_rejection = False
+        self._stream_metadata = None
         if self.runtime._credential_runner is None:
             self._call = self.stub.PollWorkerTask()
         else:
@@ -398,6 +433,7 @@ class PullStream:
             )
             if self.runtime.is_stopping:
                 return False
+            self._stream_metadata = metadata
             self._call = self.stub.PollWorkerTask(metadata=metadata)
         await self._send(
             pb.PollWorkerTaskRequest(