dagster-postgres 0.29.2__tar.gz → 0.29.3__tar.gz

{dagster_postgres-0.29.2 → dagster_postgres-0.29.3}/.gitignore

@@ -27,6 +27,7 @@ wheels/
 *.egg
 
 !docs/docs/guides/build
+!examples/docs_snippets/docs_snippets/guides/build
 
 # PyInstaller
 #  Usually these files are written by a python script from a template

{dagster_postgres-0.29.2 → dagster_postgres-0.29.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dagster-postgres
-Version: 0.29.2
+Version: 0.29.3
 Summary: A Dagster integration for postgres
 Project-URL: Homepage, https://github.com/dagster-io/dagster/tree/master/python_modules/libraries/dagster-postgres
 Author-email: Dagster Labs <hello@dagsterlabs.com>
@@ -13,6 +13,16 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
 Requires-Python: <3.15,>=3.10
-Requires-Dist: dagster==1.13.2
+Requires-Dist: dagster==1.13.3
 Requires-Dist: psycopg2-binary
+Provides-Extra: aws
+Requires-Dist: boto3>=1.26.0; extra == 'aws'
+Provides-Extra: azure
+Requires-Dist: azure-identity>=1.12.0; extra == 'azure'
+Provides-Extra: gcp
+Requires-Dist: google-auth>=2.0.0; extra == 'gcp'
 Provides-Extra: test
+Provides-Extra: wif
+Requires-Dist: azure-identity>=1.12.0; extra == 'wif'
+Requires-Dist: boto3>=1.26.0; extra == 'wif'
+Requires-Dist: google-auth>=2.0.0; extra == 'wif'

dagster_postgres-0.29.3/dagster_postgres/auth.py

@@ -0,0 +1,195 @@
+import abc
+import logging
+import threading
+import time
+from typing import Any
+
+import sqlalchemy
+from sqlalchemy import event
+
+logger = logging.getLogger(__name__)
+
+
+class PgTokenProvider(abc.ABC):
+    """Abstract base for providers that supply short-lived database access tokens.
+
+    Subclasses implement ``_fetch_token`` to obtain a fresh token from their
+    respective cloud provider.  Caching and thread-safety are handled by the
+    base class so that callers can simply invoke ``get_token()``.
+    """
+
+    def __init__(self) -> None:
+        self._cached_token: str | None = None
+        self._token_expiry: float = 0.0  # epoch seconds
+        self._lock = threading.Lock()
+
+    @abc.abstractmethod
+    def _fetch_token(self) -> tuple[str, float]:
+        """Return ``(token, expiry_epoch_seconds)``.
+
+        If the exact expiry is unknown, implementations should return
+        ``time.time() + 3600`` as a conservative default (1 hour).
+        """
+        ...
+
+    def get_token(self) -> str:
+        """Return a valid token, refreshing if expired or about to expire."""
+        # Refresh 5 minutes before expiry to avoid races
+        with self._lock:
+            if self._cached_token is None or time.time() >= (self._token_expiry - 300):
+                logger.debug("Refreshing PostgreSQL auth token")
+                self._cached_token, self._token_expiry = self._fetch_token()
+            return self._cached_token
+
+
+class AzureWifTokenProvider(PgTokenProvider):
+    """Azure Workload Identity Federation token provider.
+
+    Uses ``DefaultAzureCredential`` which automatically picks up the
+    ``AZURE_CLIENT_ID``, ``AZURE_TENANT_ID``, and
+    ``AZURE_FEDERATED_TOKEN_FILE`` environment variables set by the Azure
+    Workload Identity webhook on AKS pods.
+    """
+
+    def __init__(
+        self,
+        scope: str = "https://ossrdbms-aad.database.windows.net/.default",
+    ) -> None:
+        super().__init__()
+        self._scope = scope
+        self._credential: Any = None
+
+    def _fetch_token(self) -> tuple[str, float]:
+        if self._credential is None:
+            try:
+                from azure.identity import DefaultAzureCredential  # type: ignore[import-untyped]
+            except ImportError:
+                raise ImportError(
+                    "azure-identity is required for Azure WIF auth. "
+                    "Install with: pip install dagster-postgres[azure]"
+                )
+            self._credential = DefaultAzureCredential()
+        token = self._credential.get_token(self._scope)
+        return token.token, token.expires_on
+
+
+class GcpWifTokenProvider(PgTokenProvider):
+    """GCP Workload Identity Federation token provider.
+
+    Uses Application Default Credentials (ADC), which automatically picks up
+    the projected service account token on GKE pods.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._credentials: Any = None
+        self._request: Any = None
+
+    def _fetch_token(self) -> tuple[str, float]:
+        if self._credentials is None:
+            try:
+                import google.auth  # type: ignore[import-untyped]
+                import google.auth.transport.requests  # type: ignore[import-untyped]
+            except ImportError:
+                raise ImportError(
+                    "google-auth is required for GCP WIF auth. "
+                    "Install with: pip install dagster-postgres[gcp]"
+                )
+            self._credentials, _ = google.auth.default(
+                scopes=["https://www.googleapis.com/auth/cloud-platform"],
+            )
+            self._request = google.auth.transport.requests.Request()
+        self._credentials.refresh(self._request)
+        expiry = (
+            self._credentials.expiry.timestamp()
+            if self._credentials.expiry
+            else (time.time() + 3600)
+        )
+        return self._credentials.token, expiry
+
+
+class AwsWifTokenProvider(PgTokenProvider):
+    """AWS IAM / IRSA token provider for RDS.
+
+    Uses ``boto3`` to generate an RDS auth token.  Requires the hostname,
+    port, username, and AWS region from the connection config.
+    """
+
+    def __init__(
+        self,
+        hostname: str,
+        port: int,
+        username: str,
+        region: str,
+    ) -> None:
+        super().__init__()
+        self._hostname = hostname
+        self._port = port
+        self._username = username
+        self._region = region
+        self._client: Any = None
+
+    def _fetch_token(self) -> tuple[str, float]:
+        if self._client is None:
+            try:
+                import boto3  # type: ignore[import-untyped]
+            except ImportError:
+                raise ImportError(
+                    "boto3 is required for AWS WIF auth. "
+                    "Install with: pip install dagster-postgres[aws]"
+                )
+            self._client = boto3.client("rds", region_name=self._region)
+        token = self._client.generate_db_auth_token(
+            DBHostname=self._hostname,
+            Port=self._port,
+            DBUsername=self._username,
+        )
+        # RDS auth tokens are valid for 15 minutes
+        return token, time.time() + 900
+
+
+def create_token_provider(
+    auth_config: dict[str, Any],
+    db_config: dict[str, Any] | None = None,
+) -> PgTokenProvider:
+    """Create a token provider from the ``auth_provider`` config block."""
+    if "azure_wif" in auth_config:
+        azure_cfg = auth_config["azure_wif"]
+        return AzureWifTokenProvider(
+            scope=azure_cfg.get("scope", "https://ossrdbms-aad.database.windows.net/.default"),
+        )
+    elif "gcp_wif" in auth_config:
+        return GcpWifTokenProvider()
+    elif "aws_wif" in auth_config:
+        aws_cfg = auth_config["aws_wif"]
+        if db_config is None:
+            raise ValueError(
+                "AWS WIF requires postgres_db config (not postgres_url) to"
+                " extract hostname/port/username"
+            )
+        return AwsWifTokenProvider(
+            hostname=db_config["hostname"],
+            port=db_config.get("port", 5432),
+            username=db_config["username"],
+            region=aws_cfg["region"],
+        )
+    else:
+        raise ValueError(f"Unknown auth_provider config: {auth_config}")
+
+
+def register_do_connect_hook(
+    engine: sqlalchemy.engine.Engine,
+    token_provider: PgTokenProvider,
+) -> None:
+    """Register a ``do_connect`` event listener that injects a fresh token as
+    the password before each DBAPI connection.
+    """
+
+    @event.listens_for(engine, "do_connect")
+    def _inject_token(
+        dialect: Any,
+        conn_rec: Any,
+        cargs: list[Any],
+        cparams: dict[str, Any],
+    ) -> None:
+        cparams["password"] = token_provider.get_token()

{dagster_postgres-0.29.2 → dagster_postgres-0.29.3}/dagster_postgres/event_log/event_log.py

@@ -1,6 +1,6 @@
 from collections.abc import Iterator, Mapping, Sequence
 from contextlib import contextmanager
-from typing import Any, ContextManager, cast  # noqa: UP035
+from typing import TYPE_CHECKING, Any, ContextManager, cast  # noqa: UP035
 
 import dagster._check as check
 import sqlalchemy as db
@@ -36,6 +36,8 @@ from sqlalchemy.engine import Connection
 
 from dagster_postgres.utils import (
     create_pg_connection,
+    create_pg_engine,
+    get_token_provider_from_config,
     pg_alembic_config,
     pg_url_from_config,
     retry_pg_connection_fn,
@@ -43,6 +45,9 @@ from dagster_postgres.utils import (
     set_pg_statement_timeout,
 )
 
+if TYPE_CHECKING:
+    from dagster_postgres.auth import PgTokenProvider
+
 CHANNEL_NAME = "run_events"
 
 
@@ -56,7 +61,7 @@ class PostgresEventLogStorage(SqlEventLogStorage, ConfigurableClass):
     To use Postgres for all of the components of your instance storage, you can add the following
     block to your ``dagster.yaml``:
 
-    .. literalinclude:: ../../../../../../examples/docs_snippets/docs_snippets/deploying/dagster-pg.yaml
+    .. literalinclude:: ../../../../../../examples/docs_snippets/docs_snippets/deployment/oss/dagster-pg.yaml
        :caption: dagster.yaml
        :lines: 1-8
        :language: YAML
@@ -65,7 +70,7 @@ class PostgresEventLogStorage(SqlEventLogStorage, ConfigurableClass):
     configuring your event log storage to use Postgres, you can add a block such as the following
     to your ``dagster.yaml``:
 
-    .. literalinclude:: ../../../../../../examples/docs_snippets/docs_snippets/deploying/dagster-pg-legacy.yaml
+    .. literalinclude:: ../../../../../../examples/docs_snippets/docs_snippets/deployment/oss/dagster-pg-legacy.yaml
        :caption: dagster.yaml
        :lines: 12-21
        :language: YAML
@@ -80,16 +85,21 @@ class PostgresEventLogStorage(SqlEventLogStorage, ConfigurableClass):
         postgres_url: str,
         should_autocreate_tables: bool = True,
         inst_data: ConfigurableClassData | None = None,
+        token_provider: "PgTokenProvider | None" = None,
     ):
         self._inst_data = check.opt_inst_param(inst_data, "inst_data", ConfigurableClassData)
         self.postgres_url = check.str_param(postgres_url, "postgres_url")
         self.should_autocreate_tables = check.bool_param(
             should_autocreate_tables, "should_autocreate_tables"
         )
+        self._token_provider = token_provider
 
         # Default to not holding any connections open to prevent accumulating connections per DagsterInstance
-        self._engine = create_engine(
-            self.postgres_url, isolation_level="AUTOCOMMIT", poolclass=db_pool.NullPool
+        self._engine = create_pg_engine(
+            self.postgres_url,
+            self._token_provider,
+            isolation_level="AUTOCOMMIT",
+            poolclass=db_pool.NullPool,
         )
         self._event_watcher: SqlPollingEventWatcher | None = None
 
@@ -125,7 +135,7 @@ class PostgresEventLogStorage(SqlEventLogStorage, ConfigurableClass):
         existing_options = self._engine.url.query.get("options")
         if existing_options:
             kwargs["connect_args"] = {"options": existing_options}
-        self._engine = create_engine(self.postgres_url, **kwargs)
+        self._engine = create_pg_engine(self.postgres_url, self._token_provider, **kwargs)
         event.listen(
             self._engine,
             "connect",
@@ -153,6 +163,7 @@ class PostgresEventLogStorage(SqlEventLogStorage, ConfigurableClass):
             inst_data=inst_data,
             postgres_url=pg_url_from_config(config_value),
             should_autocreate_tables=config_value.get("should_autocreate_tables", True),
+            token_provider=get_token_provider_from_config(config_value),
         )
 
     @staticmethod

{dagster_postgres-0.29.2 → dagster_postgres-0.29.3}/dagster_postgres/run_storage/run_storage.py

@@ -1,6 +1,6 @@
 import zlib
 from collections.abc import Mapping
-from typing import ContextManager  # noqa: UP035
+from typing import TYPE_CHECKING, ContextManager  # noqa: UP035
 
 import dagster._check as check
 import sqlalchemy as db
@@ -31,6 +31,8 @@ from sqlalchemy.engine import Connection
 
 from dagster_postgres.utils import (
     create_pg_connection,
+    create_pg_engine,
+    get_token_provider_from_config,
     pg_alembic_config,
     pg_url_from_config,
     retry_pg_connection_fn,
@@ -38,6 +40,9 @@ from dagster_postgres.utils import (
     set_pg_statement_timeout,
 )
 
+if TYPE_CHECKING:
+    from dagster_postgres.auth import PgTokenProvider
+
 
 class PostgresRunStorage(SqlRunStorage, ConfigurableClass):
     """Postgres-backed run storage.
@@ -49,7 +54,7 @@ class PostgresRunStorage(SqlRunStorage, ConfigurableClass):
     To use Postgres for all of the components of your instance storage, you can add the following
     block to your ``dagster.yaml``:
 
-    .. literalinclude:: ../../../../../../examples/docs_snippets/docs_snippets/deploying/dagster-pg.yaml
+    .. literalinclude:: ../../../../../../examples/docs_snippets/docs_snippets/deployment/oss/dagster-pg.yaml
        :caption: dagster.yaml
        :lines: 1-8
        :language: YAML
@@ -58,7 +63,7 @@ class PostgresRunStorage(SqlRunStorage, ConfigurableClass):
     configuring your run storage to use Postgres, you can add a block such as the following
     to your ``dagster.yaml``:
 
-    .. literalinclude:: ../../../../../../examples/docs_snippets/docs_snippets/deploying/dagster-pg-legacy.yaml
+    .. literalinclude:: ../../../../../../examples/docs_snippets/docs_snippets/deployment/oss/dagster-pg-legacy.yaml
        :caption: dagster.yaml
        :lines: 1-10
        :language: YAML
@@ -72,16 +77,19 @@ class PostgresRunStorage(SqlRunStorage, ConfigurableClass):
         postgres_url: str,
         should_autocreate_tables: bool = True,
         inst_data: ConfigurableClassData | None = None,
+        token_provider: "PgTokenProvider | None" = None,
     ):
         self._inst_data = check.opt_inst_param(inst_data, "inst_data", ConfigurableClassData)
         self.postgres_url = postgres_url
         self.should_autocreate_tables = check.bool_param(
             should_autocreate_tables, "should_autocreate_tables"
         )
+        self._token_provider = token_provider
 
         # Default to not holding any connections open to prevent accumulating connections per DagsterInstance
-        self._engine = create_engine(
+        self._engine = create_pg_engine(
             self.postgres_url,
+            self._token_provider,
             isolation_level="AUTOCOMMIT",
             poolclass=db_pool.NullPool,
         )
@@ -121,7 +129,7 @@ class PostgresRunStorage(SqlRunStorage, ConfigurableClass):
         existing_options = self._engine.url.query.get("options")
         if existing_options:
             kwargs["connect_args"] = {"options": existing_options}
-        self._engine = create_engine(self.postgres_url, **kwargs)
+        self._engine = create_pg_engine(self.postgres_url, self._token_provider, **kwargs)
         event.listen(
             self._engine,
             "connect",
@@ -144,6 +152,7 @@ class PostgresRunStorage(SqlRunStorage, ConfigurableClass):
             inst_data=inst_data,
             postgres_url=pg_url_from_config(config_value),
             should_autocreate_tables=config_value.get("should_autocreate_tables", True),
+            token_provider=get_token_provider_from_config(config_value),
         )
 
     @staticmethod

{dagster_postgres-0.29.2 → dagster_postgres-0.29.3}/dagster_postgres/schedule_storage/schedule_storage.py

@@ -1,5 +1,5 @@
 from collections.abc import Sequence
-from typing import ContextManager  # noqa: UP035
+from typing import TYPE_CHECKING, ContextManager  # noqa: UP035
 
 import dagster._check as check
 import sqlalchemy as db
@@ -31,6 +31,8 @@ from sqlalchemy.engine import Connection
 
 from dagster_postgres.utils import (
     create_pg_connection,
+    create_pg_engine,
+    get_token_provider_from_config,
     pg_alembic_config,
     pg_url_from_config,
     retry_pg_connection_fn,
@@ -38,6 +40,9 @@ from dagster_postgres.utils import (
     set_pg_statement_timeout,
 )
 
+if TYPE_CHECKING:
+    from dagster_postgres.auth import PgTokenProvider
+
 
 class PostgresScheduleStorage(SqlScheduleStorage, ConfigurableClass):
     """Postgres-backed run storage.
@@ -49,7 +54,7 @@ class PostgresScheduleStorage(SqlScheduleStorage, ConfigurableClass):
     To use Postgres for all of the components of your instance storage, you can add the following
     block to your ``dagster.yaml``:
 
-    .. literalinclude:: ../../../../../../examples/docs_snippets/docs_snippets/deploying/dagster-pg.yaml
+    .. literalinclude:: ../../../../../../examples/docs_snippets/docs_snippets/deployment/oss/dagster-pg.yaml
        :caption: dagster.yaml
        :lines: 1-8
        :language: YAML
@@ -58,7 +63,7 @@ class PostgresScheduleStorage(SqlScheduleStorage, ConfigurableClass):
     configuring your schedule storage to use Postgres, you can add a block such as the following
     to your ``dagster.yaml``:
 
-    .. literalinclude:: ../../../../../../examples/docs_snippets/docs_snippets/deploying/dagster-pg-legacy.yaml
+    .. literalinclude:: ../../../../../../examples/docs_snippets/docs_snippets/deployment/oss/dagster-pg-legacy.yaml
        :caption: dagster.yaml
        :lines: 23-32
        :language: YAML
@@ -72,16 +77,21 @@ class PostgresScheduleStorage(SqlScheduleStorage, ConfigurableClass):
         postgres_url: str,
         should_autocreate_tables: bool = True,
         inst_data: ConfigurableClassData | None = None,
+        token_provider: "PgTokenProvider | None" = None,
     ):
         self._inst_data = check.opt_inst_param(inst_data, "inst_data", ConfigurableClassData)
         self.postgres_url = postgres_url
         self.should_autocreate_tables = check.bool_param(
             should_autocreate_tables, "should_autocreate_tables"
         )
+        self._token_provider = token_provider
 
         # Default to not holding any connections open to prevent accumulating connections per DagsterInstance
-        self._engine = create_engine(
-            self.postgres_url, isolation_level="AUTOCOMMIT", poolclass=db_pool.NullPool
+        self._engine = create_pg_engine(
+            self.postgres_url,
+            self._token_provider,
+            isolation_level="AUTOCOMMIT",
+            poolclass=db_pool.NullPool,
         )
 
         # Stamp and create tables if the main table does not exist (we can't check alembic
@@ -117,7 +127,7 @@ class PostgresScheduleStorage(SqlScheduleStorage, ConfigurableClass):
         existing_options = self._engine.url.query.get("options")
         if existing_options:
             kwargs["connect_args"] = {"options": existing_options}
-        self._engine = create_engine(self.postgres_url, **kwargs)
+        self._engine = create_pg_engine(self.postgres_url, self._token_provider, **kwargs)
         event.listen(
             self._engine,
             "connect",
@@ -140,6 +150,7 @@ class PostgresScheduleStorage(SqlScheduleStorage, ConfigurableClass):
             inst_data=inst_data,
             postgres_url=pg_url_from_config(config_value),
             should_autocreate_tables=config_value.get("should_autocreate_tables", True),
+            token_provider=get_token_provider_from_config(config_value),
         )
 
     @staticmethod

{dagster_postgres-0.29.2 → dagster_postgres-0.29.3}/dagster_postgres/storage.py

@@ -1,3 +1,5 @@
+from typing import TYPE_CHECKING
+
 from dagster import _check as check
 from dagster._config.config_schema import UserConfigSchema
 from dagster._core.storage.base_storage import DagsterStorage
@@ -10,7 +12,10 @@ from dagster._serdes import ConfigurableClass, ConfigurableClassData
 from dagster_postgres.event_log import PostgresEventLogStorage
 from dagster_postgres.run_storage import PostgresRunStorage
 from dagster_postgres.schedule_storage import PostgresScheduleStorage
-from dagster_postgres.utils import pg_url_from_config
+from dagster_postgres.utils import get_token_provider_from_config, pg_url_from_config
+
+if TYPE_CHECKING:
+    from dagster_postgres.auth import PgTokenProvider
 
 
 class DagsterPostgresStorage(DagsterStorage, ConfigurableClass):
@@ -23,7 +28,7 @@ class DagsterPostgresStorage(DagsterStorage, ConfigurableClass):
     To use Postgres for storage, you can add a block such as the following to your
     ``dagster.yaml``:
 
-    .. literalinclude:: ../../../../../../examples/docs_snippets/docs_snippets/deploying/dagster-pg.yaml
+    .. literalinclude:: ../../../../../../examples/docs_snippets/docs_snippets/deployment/oss/dagster-pg.yaml
        :caption: dagster.yaml
        :lines: 1-8
        :language: YAML
@@ -34,18 +39,25 @@ class DagsterPostgresStorage(DagsterStorage, ConfigurableClass):
 
     def __init__(
         self,
-        postgres_url,
-        should_autocreate_tables=True,
+        postgres_url: str,
+        should_autocreate_tables: bool = True,
         inst_data: ConfigurableClassData | None = None,
+        token_provider: "PgTokenProvider | None" = None,
     ):
         self.postgres_url = postgres_url
         self.should_autocreate_tables = check.bool_param(
             should_autocreate_tables, "should_autocreate_tables"
         )
         self._inst_data = check.opt_inst_param(inst_data, "inst_data", ConfigurableClassData)
-        self._run_storage = PostgresRunStorage(postgres_url, should_autocreate_tables)
-        self._event_log_storage = PostgresEventLogStorage(postgres_url, should_autocreate_tables)
-        self._schedule_storage = PostgresScheduleStorage(postgres_url, should_autocreate_tables)
+        self._run_storage = PostgresRunStorage(
+            postgres_url, should_autocreate_tables, token_provider=token_provider
+        )
+        self._event_log_storage = PostgresEventLogStorage(
+            postgres_url, should_autocreate_tables, token_provider=token_provider
+        )
+        self._schedule_storage = PostgresScheduleStorage(
+            postgres_url, should_autocreate_tables, token_provider=token_provider
+        )
         super().__init__()
 
     @property
@@ -64,6 +76,7 @@ class DagsterPostgresStorage(DagsterStorage, ConfigurableClass):
             inst_data=inst_data,
             postgres_url=pg_url_from_config(config_value),
             should_autocreate_tables=config_value.get("should_autocreate_tables", True),
+            token_provider=get_token_provider_from_config(config_value),
         )
 
     @property

{dagster_postgres-0.29.2 → dagster_postgres-0.29.3}/dagster_postgres/utils.py

@@ -2,7 +2,7 @@ import logging
 import time
 from collections.abc import Callable, Iterator, Mapping
 from contextlib import contextmanager
-from typing import Any, TypeVar
+from typing import TYPE_CHECKING, Any, TypeVar
 from urllib.parse import quote, urlencode
 
 import alembic.config
@@ -13,9 +13,12 @@ from dagster._core.definitions.policy import Backoff, Jitter, calculate_delay
 
 # re-export
 from dagster._core.storage.config import pg_config as pg_config
-from dagster._core.storage.sql import get_alembic_config
+from dagster._core.storage.sql import create_engine, get_alembic_config
 from sqlalchemy.engine import Connection
 
+if TYPE_CHECKING:
+    from dagster_postgres.auth import PgTokenProvider
+
 T = TypeVar("T")
 
 
@@ -41,14 +44,18 @@ def pg_url_from_config(config_value: Mapping[str, Any]) -> str:
 
 def get_conn_string(
     username: str,
-    password: str,
-    hostname: str,
-    db_name: str,
+    password: str = "",
+    hostname: str = "",
+    db_name: str = "",
     port: str = "5432",
     params: Mapping[str, object] | None = None,
     scheme: str = "postgresql",
 ) -> str:
-    uri = f"{scheme}://{quote(username)}:{quote(password)}@{hostname}:{port}/{db_name}"
+    if password:
+        userinfo = f"{quote(username)}:{quote(password)}"
+    else:
+        userinfo = quote(username)
+    uri = f"{scheme}://{userinfo}@{hostname}:{port}/{db_name}"
 
     if params:
         query_string = f"{urlencode(params, quote_via=quote)}"
@@ -157,3 +164,41 @@ def set_pg_statement_timeout(conn: Any, millis: int):
     with conn:
         with conn.cursor() as curs:
             curs.execute(f"SET statement_timeout = {millis};")
+
+
+def create_pg_engine(
+    postgres_url: str,
+    token_provider: "PgTokenProvider | None" = None,
+    **engine_kwargs: Any,
+) -> sqlalchemy.engine.Engine:
+    """Create a SQLAlchemy engine, optionally with WIF token injection via
+    the ``do_connect`` event hook.
+    """
+    engine = create_engine(postgres_url, **engine_kwargs)
+    if token_provider is not None:
+        from dagster_postgres.auth import register_do_connect_hook
+
+        register_do_connect_hook(engine, token_provider)
+    return engine
+
+
+def get_token_provider_from_config(
+    config_value: Mapping[str, Any],
+) -> "PgTokenProvider | None":
+    """Extract and build a token provider from config, or return ``None`` for
+    password auth.
+    """
+    auth_config = config_value.get("auth_provider")
+    if auth_config is None:
+        return None
+
+    check.invariant(
+        "postgres_url" not in config_value or not config_value["postgres_url"],
+        "auth_provider cannot be used with postgres_url. Use postgres_db instead"
+        " so that tokens can be injected at the DBAPI connection level.",
+    )
+
+    from dagster_postgres.auth import create_token_provider
+
+    db_config = config_value.get("postgres_db")
+    return create_token_provider(auth_config, db_config)

dagster_postgres-0.29.3/dagster_postgres/version.py

@@ -0,0 +1 @@
+__version__ = "0.29.3"

dagster_postgres-0.29.3/dagster_postgres_tests/test_auth.py

@@ -0,0 +1,229 @@
+import threading
+import time
+from unittest.mock import patch
+
+import pytest
+import sqlalchemy
+import sqlalchemy.exc
+import sqlalchemy.pool
+from dagster_postgres.auth import (
+    AwsWifTokenProvider,
+    AzureWifTokenProvider,
+    GcpWifTokenProvider,
+    PgTokenProvider,
+    create_token_provider,
+    register_do_connect_hook,
+)
+from dagster_postgres.utils import get_token_provider_from_config
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+class FakeTokenProvider(PgTokenProvider):
+    """A concrete token provider that returns a fixed token for testing."""
+
+    def __init__(self, token: str = "fake-token", ttl: float = 3600) -> None:
+        super().__init__()
+        self._token = token
+        self._ttl = ttl
+        self.fetch_count = 0
+
+    def _fetch_token(self) -> tuple[str, float]:
+        self.fetch_count += 1
+        return self._token, time.time() + self._ttl
+
+
+# ---------------------------------------------------------------------------
+# Token caching
+# ---------------------------------------------------------------------------
+
+
+class TestTokenCaching:
+    def test_caches_token_within_ttl(self) -> None:
+        provider = FakeTokenProvider(token="tok-1", ttl=3600)
+        assert provider.get_token() == "tok-1"
+        assert provider.get_token() == "tok-1"
+        assert provider.fetch_count == 1
+
+    def test_refreshes_on_expiry(self) -> None:
+        provider = FakeTokenProvider(token="tok-1", ttl=0)
+        # ttl=0 means the token is already expired on creation
+        tok1 = provider.get_token()
+        assert tok1 == "tok-1"
+        assert provider.fetch_count == 1
+
+        # Second call should refresh because token is already past expiry
+        tok2 = provider.get_token()
+        assert tok2 == "tok-1"
+        assert provider.fetch_count == 2
+
+    def test_refreshes_within_margin(self) -> None:
+        # Token expires in 4 minutes — within the 5-minute refresh margin
+        provider = FakeTokenProvider(token="tok-1", ttl=240)
+        provider.get_token()
+        assert provider.fetch_count == 1
+
+        # Should refresh because 240s < 300s margin
+        provider.get_token()
+        assert provider.fetch_count == 2
+
+
+# ---------------------------------------------------------------------------
+# Thread safety
+# ---------------------------------------------------------------------------
+
+
+class TestThreadSafety:
+    def test_concurrent_get_token(self) -> None:
+        provider = FakeTokenProvider(token="tok-1", ttl=3600)
+        results: list[str] = []
+        errors: list[Exception] = []
+
+        def worker() -> None:
+            try:
+                results.append(provider.get_token())
+            except Exception as e:
+                errors.append(e)
+
+        threads = [threading.Thread(target=worker) for _ in range(20)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        assert not errors
+        assert all(r == "tok-1" for r in results)
+        # May have fetched more than once due to race at startup, but should be bounded
+        assert provider.fetch_count >= 1
+
+
+# ---------------------------------------------------------------------------
+# do_connect hook
+# ---------------------------------------------------------------------------
+
+
+class TestDoConnectHook:
+    def test_hook_injects_token_on_connect(self) -> None:
+        """Verify that register_do_connect_hook causes the provider's token
+        to be used as the password when psycopg2.connect is called.
+        """
+        provider = FakeTokenProvider(token="injected-password", ttl=3600)
+        engine = sqlalchemy.create_engine(
+            "postgresql://user:dummy@localhost:5432/db",
+            poolclass=sqlalchemy.pool.NullPool,
+        )
+        register_do_connect_hook(engine, provider)
+
+        # Patch psycopg2.connect to capture the password that SQLAlchemy passes
+        captured_kwargs: dict[str, object] = {}
+        with patch("psycopg2.connect") as mock_connect:
+
+            def capture_connect(*args: object, **kwargs: object) -> None:
+                captured_kwargs.update(kwargs)
+                raise sqlalchemy.exc.OperationalError("test", {}, Exception("skip"))
+
+            mock_connect.side_effect = capture_connect
+
+            # Attempting to connect will invoke do_connect, which injects the token,
+            # then psycopg2.connect is called with the modified password.
+            try:
+                with engine.connect():
+                    pass
+            except Exception:
+                pass
+
+        # The do_connect hook should have replaced "dummy" with the token
+        assert captured_kwargs.get("password") == "injected-password"
+        assert provider.fetch_count == 1
+
+
+# ---------------------------------------------------------------------------
+# Factory function
+# ---------------------------------------------------------------------------
+
+
+class TestCreateTokenProvider:
+    def test_azure_wif(self) -> None:
+        provider = create_token_provider({"azure_wif": {"scope": "custom-scope"}})
+        assert isinstance(provider, AzureWifTokenProvider)
+        assert provider._scope == "custom-scope"  # noqa: SLF001
+
+    def test_azure_wif_default_scope(self) -> None:
+        provider = create_token_provider({"azure_wif": {}})
+        assert isinstance(provider, AzureWifTokenProvider)
+        assert "ossrdbms-aad" in provider._scope  # noqa: SLF001
+
+    def test_gcp_wif(self) -> None:
+        provider = create_token_provider({"gcp_wif": {}})
+        assert isinstance(provider, GcpWifTokenProvider)
+
+    def test_aws_wif(self) -> None:
+        db_config = {"hostname": "my-host", "port": 5432, "username": "myuser"}
+        provider = create_token_provider({"aws_wif": {"region": "us-east-1"}}, db_config)
+        assert isinstance(provider, AwsWifTokenProvider)
+        assert provider._region == "us-east-1"  # noqa: SLF001
+        assert provider._hostname == "my-host"  # noqa: SLF001
+
+    def test_aws_wif_requires_db_config(self) -> None:
+        with pytest.raises(ValueError, match="AWS WIF requires postgres_db"):
+            create_token_provider({"aws_wif": {"region": "us-east-1"}})
+
+    def test_unknown_provider(self) -> None:
+        with pytest.raises(ValueError, match="Unknown auth_provider"):
+            create_token_provider({"unknown_wif": {}})
+
+
+# ---------------------------------------------------------------------------
+# Config extraction
+# ---------------------------------------------------------------------------
+
+
+class TestGetTokenProviderFromConfig:
+    def test_returns_none_without_auth_provider(self) -> None:
+        config = {
+            "postgres_db": {"username": "u", "password": "p", "hostname": "h", "db_name": "d"}
+        }
+        assert get_token_provider_from_config(config) is None
+
+    def test_returns_provider_with_auth_provider(self) -> None:
+        config = {
+            "postgres_db": {"username": "u", "hostname": "h", "db_name": "d"},
+            "auth_provider": {"azure_wif": {}},
+        }
+        provider = get_token_provider_from_config(config)
+        assert isinstance(provider, AzureWifTokenProvider)
+
+    def test_rejects_auth_provider_with_postgres_url(self) -> None:
+        config = {
+            "postgres_url": "postgresql://user:pass@host/db",
+            "auth_provider": {"azure_wif": {}},
+        }
+        with pytest.raises(Exception, match="auth_provider cannot be used with postgres_url"):
+            get_token_provider_from_config(config)
+
+
+# ---------------------------------------------------------------------------
+# Cloud provider import errors
+# ---------------------------------------------------------------------------
+
+
+class TestImportErrors:
+    def test_azure_missing_import(self) -> None:
+        provider = AzureWifTokenProvider()
+        with patch.dict("sys.modules", {"azure.identity": None, "azure": None}):
+            with pytest.raises(ImportError, match="azure-identity"):
+                provider._fetch_token()  # noqa: SLF001
+
+    def test_gcp_missing_import(self) -> None:
+        provider = GcpWifTokenProvider()
+        with patch.dict("sys.modules", {"google.auth": None, "google": None}):
+            with pytest.raises(ImportError, match="google-auth"):
+                provider._fetch_token()  # noqa: SLF001
+
+    def test_aws_missing_import(self) -> None:
+        provider = AwsWifTokenProvider(hostname="h", port=5432, username="u", region="us-east-1")
+        with patch.dict("sys.modules", {"boto3": None}):
+            with pytest.raises(ImportError, match="boto3"):
+                provider._fetch_token()  # noqa: SLF001

{dagster_postgres-0.29.2 → dagster_postgres-0.29.3}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "dagster-postgres"
-version = "0.29.2"
+version = "0.29.3"
 description = "A Dagster integration for postgres"
 authors = [
     { name = "Dagster Labs", email = "hello@dagsterlabs.com" },
@@ -22,11 +22,25 @@ classifiers = [
     "Operating System :: OS Independent",
 ]
 dependencies = [
-    "dagster==1.13.2",
+    "dagster==1.13.3",
     "psycopg2-binary",
 ]
 
 [project.optional-dependencies]
+azure = [
+    "azure-identity>=1.12.0",
+]
+gcp = [
+    "google-auth>=2.0.0",
+]
+aws = [
+    "boto3>=1.26.0",
+]
+wif = [
+    "azure-identity>=1.12.0",
+    "google-auth>=2.0.0",
+    "boto3>=1.26.0",
+]
 test = []
 
 [project.urls]

dagster_postgres-0.29.2/dagster_postgres/version.py

@@ -1 +0,0 @@
-__version__ = "0.29.2"