feat(postgresql): add Azure Entra CLI authentication (#317)

Co-authored-by: Peter Adams <18162810+Maxteabag@users.noreply.github.com>

Author: Maxteabag

sqlit/cli.py

@@ -727,6 +727,11 @@ def main() -> int:
         choices=[t.value for t in AuthType],
         help="Authentication type (SQL Server only)",
     )
+    edit_parser.add_argument(
+        "--postgres-auth-method",
+        choices=["password", "azure_entra_cli"],
+        help="PostgreSQL authentication method",
+    )
     edit_parser.add_argument("--file-path", help="Database file path (SQLite only)")
     edit_parser.add_argument("--password-command", dest="password_command", help="Shell command to retrieve the database password")
     edit_parser.add_argument("--ssh-password-command", dest="ssh_password_command", help="Shell command to retrieve the SSH password")

sqlit/core/connection_manager.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import copy
 from dataclasses import dataclass
 from typing import Any
 
@@ -32,19 +33,38 @@ def populate_credentials(self, config: ConnectionConfig) -> ConnectionConfig:
             return config
 
         service = self._services.credentials_service
-        if endpoint and endpoint.password is None and uses_db_password(config):
+        if endpoint and endpoint.password is None and not endpoint.password_command and uses_db_password(config):
             password = service.get_password(config.name)
             if password is not None:
                 endpoint.password = password
-        if config.tunnel and config.tunnel.password is None:
+        if config.tunnel and config.tunnel.password is None and not config.tunnel.password_command:
             ssh_password = service.get_ssh_password(config.name)
             if ssh_password is not None:
                 config.tunnel.password = ssh_password
         return config
 
+    def _resolve_dynamic_credentials(self, config: ConnectionConfig) -> ConnectionConfig:
+        """Resolve commands on a copy so short-lived secrets are not persisted."""
+        from sqlit.domains.connections.domain.password_command import (
+            run_password_command,
+        )
+        from sqlit.domains.connections.providers.config_service import (
+            normalize_connection_config,
+        )
+
+        resolved = normalize_connection_config(copy.deepcopy(config))
+        endpoint = resolved.tcp_endpoint
+        if endpoint and endpoint.password is None and endpoint.password_command and resolved.db_type != "postgresql":
+            # PostgreSQL resolves this at the adapter boundary so reconnects and
+            # database switches always fetch a fresh short-lived token.
+            endpoint.password = run_password_command(endpoint.password_command)
+        if resolved.tunnel and resolved.tunnel.auth_type == "password" and resolved.tunnel.password is None and resolved.tunnel.password_command:
+            resolved.tunnel.password = run_password_command(resolved.tunnel.password_command)
+        return resolved
+
     def connect(self, config: ConnectionConfig) -> Any:
         """Create a session for the given config."""
-        return self._services.session_factory(config)
+        return self._services.session_factory(self._resolve_dynamic_credentials(config))
 
     def test_connection(self, config: ConnectionConfig) -> ConnectionTestResult:
         """Test a connection without mutating UI state."""
@@ -55,13 +75,14 @@ def test_connection(self, config: ConnectionConfig) -> ConnectionTestResult:
         error: Exception | None = None
 
         try:
-            tunnel, host, port = self._services.tunnel_factory(config)
+            resolved = self._resolve_dynamic_credentials(config)
+            tunnel, host, port = self._services.tunnel_factory(resolved)
             if tunnel:
-                connect_config = config.with_endpoint(host=host, port=str(port))
+                connect_config = resolved.with_endpoint(host=host, port=str(port))
             else:
-                connect_config = config
+                connect_config = resolved
 
-            provider = self._services.provider_factory(config.db_type)
+            provider = self._services.provider_factory(resolved.db_type)
             conn = provider.connection_factory.connect(connect_config)
             conn.close()
         except Exception as exc:

sqlit/domains/connections/app/persist_utils.py

@@ -22,12 +22,16 @@ def build_persist_connections(
     persist_connections = copy.deepcopy(connections)
     for conn in persist_connections:
         endpoint = conn.tcp_endpoint
-        if endpoint and endpoint.password is None:
+        if endpoint and endpoint.password is None and not endpoint.password_command:
             stored = credentials_service.get_password(conn.name)
             if stored is not None:
                 endpoint.password = stored
 
-        if conn.tunnel and conn.tunnel.password is None:
+        if (
+            conn.tunnel
+            and conn.tunnel.password is None
+            and not conn.tunnel.password_command
+        ):
             stored_ssh = credentials_service.get_ssh_password(conn.name)
             if stored_ssh is not None:
                 conn.tunnel.password = stored_ssh

sqlit/domains/connections/cli/commands.py

@@ -277,6 +277,12 @@ def cmd_connection_edit(args: Any, *, services: AppServices | None = None) -> in
             valid_types = ", ".join(t.value for t in AuthType)
             print(f"Error: Invalid auth type '{args.auth_type}'. Valid types: {valid_types}")
             return 1
+    postgres_auth_method = getattr(args, "postgres_auth_method", None)
+    if postgres_auth_method is not None:
+        if conn.db_type != "postgresql":
+            print("Error: --postgres-auth-method is only valid for PostgreSQL.")
+            return 1
+        conn.set_option("postgres_auth_method", postgres_auth_method)
     if endpoint:
         if args.username is not None:
             endpoint.username = args.username
@@ -286,6 +292,24 @@ def cmd_connection_edit(args: Any, *, services: AppServices | None = None) -> in
     password_command = getattr(args, "password_command", None)
     if password_command is not None and endpoint:
         endpoint.password_command = password_command or None
+        if conn.db_type == "postgresql":
+            from sqlit.domains.connections.providers.postgresql.auth import (
+                AZURE_ENTRA_PASSWORD_COMMAND,
+                POSTGRES_AUTH_AZURE_ENTRA_CLI,
+            )
+
+            if password_command == AZURE_ENTRA_PASSWORD_COMMAND:
+                conn.set_option("postgres_auth_method", POSTGRES_AUTH_AZURE_ENTRA_CLI)
+                postgres_auth_method = POSTGRES_AUTH_AZURE_ENTRA_CLI
+
+    if conn.db_type == "postgresql":
+        from sqlit.domains.connections.providers.config_service import (
+            normalize_connection_config,
+        )
+
+        conn = normalize_connection_config(conn)
+        connections[conn_idx] = conn
+        endpoint = conn.tcp_endpoint
 
     ssh_password_command = getattr(args, "ssh_password_command", None)
     if ssh_password_command is not None and conn.tunnel:

sqlit/domains/connections/cli/helpers.py

@@ -15,11 +15,7 @@
 def _get_connection_arg_names() -> set[str]:
     from sqlit.domains.connections.providers.catalog import get_all_schemas
 
-    names = {
-        field.name
-        for schema in get_all_schemas().values()
-        for field in schema.fields
-    }
+    names = {field.name for schema in get_all_schemas().values() for field in schema.fields}
     names.add("host")
     return names
 
@@ -119,6 +115,14 @@ def build_connection_config_from_args(
     password_command = getattr(args, "password_command", None)
     if password_command:
         config_values["password_command"] = password_command
+        if schema.db_type == "postgresql":
+            from sqlit.domains.connections.providers.postgresql.auth import (
+                AZURE_ENTRA_PASSWORD_COMMAND,
+                POSTGRES_AUTH_AZURE_ENTRA_CLI,
+            )
+
+            if password_command == AZURE_ENTRA_PASSWORD_COMMAND:
+                config_values["postgres_auth_method"] = POSTGRES_AUTH_AZURE_ENTRA_CLI
     ssh_password_command = getattr(args, "ssh_password_command", None)
     if ssh_password_command:
         config_values["ssh_password_command"] = ssh_password_command

sqlit/domains/connections/domain/passwords.py

@@ -18,6 +18,15 @@ def uses_db_password(config: ConnectionConfig) -> bool:
     if auth_type in ("ad_default", "ad_integrated", "windows"):
         return False
 
+    if config.db_type == "postgresql":
+        from sqlit.domains.connections.providers.postgresql.auth import (
+            POSTGRES_AUTH_AZURE_ENTRA_CLI,
+            get_postgres_auth_method,
+        )
+
+        if get_postgres_auth_method(config) == POSTGRES_AUTH_AZURE_ENTRA_CLI:
+            return False
+
     if config.db_type == "trino":
         auth_method = str(config.options.get("trino_auth_method", config.extra_options.get("trino_auth_method", "basic"))).lower()
         if auth_method in {"none", "kerberos", "gssapi"}:

sqlit/domains/connections/providers/postgresql/adapter.py

@@ -72,8 +72,16 @@ def install_package(self) -> str:
     def driver_import_names(self) -> tuple[str, ...]:
         return ("psycopg2",)
 
+    def normalize_config(self, config: ConnectionConfig) -> ConnectionConfig:
+        from sqlit.domains.connections.providers.postgresql.auth import (
+            normalize_postgres_auth,
+        )
+
+        return normalize_postgres_auth(config)
+
     def connect(self, config: ConnectionConfig) -> Any:
         """Connect to PostgreSQL database."""
+        config = self.normalize_config(config)
         psycopg2 = self._import_driver_module(
             "psycopg2",
             driver_name=self.name,
@@ -100,8 +108,15 @@ def connect(self, config: ConnectionConfig) -> Any:
             connect_args["port"] = int(endpoint.port or get_default_port("postgresql"))
         if endpoint.username:
             connect_args["user"] = endpoint.username
-        if endpoint.password is not None:
-            connect_args["password"] = endpoint.password
+        password = endpoint.password
+        if password is None and endpoint.password_command:
+            from sqlit.domains.connections.domain.password_command import (
+                run_password_command,
+            )
+
+            password = run_password_command(endpoint.password_command)
+        if password is not None:
+            connect_args["password"] = password
 
         tls_mode = get_tls_mode(config)
         tls_ca, tls_cert, tls_key, tls_key_password = get_tls_files(config)

sqlit/domains/connections/providers/postgresql/auth.py

@@ -0,0 +1,32 @@
+"""Authentication helpers for PostgreSQL connections."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from sqlit.domains.connections.domain.config import ConnectionConfig
+
+POSTGRES_AUTH_PASSWORD = "password"
+POSTGRES_AUTH_AZURE_ENTRA_CLI = "azure_entra_cli"
+AZURE_ENTRA_PASSWORD_COMMAND = "az account get-access-token --resource-type oss-rdbms --query accessToken --output tsv"
+
+
+def get_postgres_auth_method(config: ConnectionConfig) -> str:
+    return str(config.get_option("postgres_auth_method", POSTGRES_AUTH_PASSWORD)).lower()
+
+
+def normalize_postgres_auth(config: ConnectionConfig) -> ConnectionConfig:
+    endpoint = config.tcp_endpoint
+    if endpoint is None:
+        return config
+    if "postgres_auth_method" not in config.options and endpoint.password_command == AZURE_ENTRA_PASSWORD_COMMAND:
+        # Preserve and adopt connections that used the generic password-command
+        # support before Azure Entra became a first-class authentication mode.
+        config.set_option("postgres_auth_method", POSTGRES_AUTH_AZURE_ENTRA_CLI)
+    if get_postgres_auth_method(config) == POSTGRES_AUTH_AZURE_ENTRA_CLI:
+        endpoint.password = None
+        endpoint.password_command = AZURE_ENTRA_PASSWORD_COMMAND
+    elif endpoint.password_command == AZURE_ENTRA_PASSWORD_COMMAND:
+        endpoint.password_command = None
+    return config

sqlit/domains/connections/providers/postgresql/schema.py

@@ -1,25 +1,56 @@
 """Connection schema for PostgreSQL."""
 
+from sqlit.domains.connections.providers.postgresql.auth import (
+    POSTGRES_AUTH_AZURE_ENTRA_CLI,
+    POSTGRES_AUTH_PASSWORD,
+)
 from sqlit.domains.connections.providers.schema_helpers import (
     SSH_FIELDS,
     TLS_FIELDS,
     ConnectionSchema,
+    FieldType,
+    SchemaField,
+    SelectOption,
     _database_field,
-    _password_field,
     _port_field,
     _server_field,
     _username_field,
 )
 
+
+def _auth_is_password(values: dict) -> bool:
+    return values.get("postgres_auth_method", POSTGRES_AUTH_PASSWORD) == POSTGRES_AUTH_PASSWORD
+
+
 SCHEMA = ConnectionSchema(
     db_type="postgresql",
     display_name="PostgreSQL",
     fields=(
         _server_field(required=False),
         _port_field("5432"),
         _database_field(),
+        SchemaField(
+            name="postgres_auth_method",
+            label="Authentication",
+            field_type=FieldType.DROPDOWN,
+            options=(
+                SelectOption(POSTGRES_AUTH_PASSWORD, "Password"),
+                SelectOption(
+                    POSTGRES_AUTH_AZURE_ENTRA_CLI,
+                    "Azure Entra ID (Azure CLI)",
+                ),
+            ),
+            default=POSTGRES_AUTH_PASSWORD,
+        ),
         _username_field(required=False),
-        _password_field(),
+        SchemaField(
+            name="password",
+            label="Password",
+            field_type=FieldType.PASSWORD,
+            placeholder="(empty = ask every connect)",
+            group="credentials",
+            visible_when=_auth_is_password,
+        ),
     )
     + SSH_FIELDS
     + TLS_FIELDS,

sqlit/domains/connections/store/connections.py

@@ -130,12 +130,16 @@ def _load_credentials(self, config: ConnectionConfig) -> None:
             config: ConnectionConfig to populate with credentials.
         """
         endpoint = config.tcp_endpoint
-        if endpoint and endpoint.password is None:
+        if endpoint and endpoint.password is None and not endpoint.password_command:
             password = self.credentials_service.get_password(config.name)
             if password is not None:
                 endpoint.password = password
 
-        if config.tunnel and config.tunnel.password is None:
+        if (
+            config.tunnel
+            and config.tunnel.password is None
+            and not config.tunnel.password_command
+        ):
             ssh_password = self.credentials_service.get_ssh_password(config.name)
             if ssh_password is not None:
                 config.tunnel.password = ssh_password
@@ -295,9 +299,17 @@ def save_one(
             target = copy.deepcopy(connection)
             try:
                 endpoint = target.tcp_endpoint
-                if endpoint and endpoint.password is None:
+                if (
+                    endpoint
+                    and endpoint.password is None
+                    and not endpoint.password_command
+                ):
                     endpoint.password = self.credentials_service.get_password_for_migration(previous_name)  # type: ignore[arg-type]
-                if target.tunnel and target.tunnel.password is None:
+                if (
+                    target.tunnel
+                    and target.tunnel.password is None
+                    and not target.tunnel.password_command
+                ):
                     target.tunnel.password = self.credentials_service.get_ssh_password_for_migration(previous_name)  # type: ignore[arg-type]
                 destination_db_password = self.credentials_service.get_password_for_migration(connection.name)
                 destination_ssh_password = self.credentials_service.get_ssh_password_for_migration(