pg-azure_workload_identity 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: 76a51f189884a5080a32166397a3cde091c99c1db2b4fc6b3e25d14ae8a24c53
+  data.tar.gz: 82c6a01050444c0a9d4d99993098dc69aef71402b64774de6ebfe0a6c8f4c452
+SHA512:
+  metadata.gz: 2c3e82ebca506d35b28d02a128e0a9a1ea6ffc21295afe28341d4291820f1a917babd3ceea8a295725328a20e39b1feddf50e4b07d8b8c890b58a90700a17c4b
+  data.tar.gz: c11bee40a5458703988050efd757f90a3e154aec5bda21604fc5d2dc3c26afe69540930fc97c2841185d71c7b226d44f8b3c0d7a0b9d470684f1f3dae1725e18

data/.yardopts

@@ -0,0 +1,6 @@
+--markup markdown
+--markup-provider redcarpet
+--no-private
+--exclude ^sig/
+-
+*.md

data/CHANGELOG.md

@@ -0,0 +1,10 @@
+# Changelog
+
+All notable changes to this project will be documented in this file.
+
+The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+## [0.1.0] - 2026-05-19
+
+### Added
+* A plugin for the [`pg` gem](https://rubygems.org/gems/pg) that adds support for [Entra ID Authentication](https://learn.microsoft.com/azure/postgresql/security/security-entra-concepts) when connecting to PostgreSQL flexible database servers in Microsoft Azure. ([#1](https://github.com/atpoint-cloud/pg-azure_workload_identity/pull/1))

data/LICENSE.txt

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2026 Simon Schmid
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

data/README.md

@@ -0,0 +1,160 @@
+# PG::AzureWorkloadIdentity
+
+[![Gem](https://img.shields.io/gem/v/pg-azure_workload_identity?style=flat-square)](https://rubygems.org/gems/pg-azure_workload_identity)
+ 
+[![Docs](https://img.shields.io/badge/yard-docs-blue?style=flat-square)](https://www.rubydoc.info/gems/pg-azure_workload_identity)
+
+`PG::AzureWorkloadIdentity` is a plugin for the [`pg` gem](https://rubygems.org/gems/pg) that adds support for [Microsoft Entra ID authentication](https://learn.microsoft.com/azure/postgresql/security/security-entra-concepts) when connecting to [Azure Database for PostgreSQL flexible server](https://learn.microsoft.com/azure/postgresql/flexible-server/overview), using [Microsoft Entra Workload Identity](https://learn.microsoft.com/entra/workload-id/workload-identity-federation) — typically as projected into pods by [AKS Workload Identity](https://learn.microsoft.com/azure/aks/workload-identity-overview).
+
+Entra ID authentication lets your application connect to the database using short-lived access tokens instead of a fixed password. Combined with Workload Identity, your pods never carry a long-lived secret at all: the Kubernetes service account JWT is automatically exchanged for an Entra access token via [federated credentials](https://learn.microsoft.com/entra/workload-id/workload-identity-federation), and that token is used as the database password. No client secrets, no certificate rotation, no password storage.
+
+## Installation
+
+Install manually:
+
+```console
+$ gem install pg-azure_workload_identity
+```
+
+or with Bundler:
+
+```console
+$ bundle add pg-azure_workload_identity
+```
+
+## Usage
+
+To use Workload Identity authentication for your database connections, you need to
+
+1. enable Microsoft Entra authentication on your Azure Database for PostgreSQL flexible server,
+2. provide your application with a Workload Identity, and
+3. configure your application to use the gem's token generator.
+
+### 1. Enable Microsoft Entra authentication on your database
+
+Configure your Azure Database for PostgreSQL flexible server to allow Microsoft Entra authentication, either during server provisioning or on an existing server through the **Security → Authentication** pane. See [Use Microsoft Entra ID for authentication with Azure Database for PostgreSQL](https://learn.microsoft.com/azure/postgresql/security/security-entra-configure) for the step-by-step instructions.
+
+Choose between **Microsoft Entra authentication only** (no password-based logins) and **PostgreSQL and Microsoft Entra authentication** (both modes available), then add at least one Microsoft Entra administrator. Only an Entra administrator can subsequently create additional Entra-mapped database roles.
+
+Connect to the database as an Entra administrator and create a role for your application's identity. For a user-assigned managed identity, use its object (principal) ID:
+
+```sql
+select * from pgaadauth_create_principal_with_oid('my-app', '<managed-identity-object-id>', 'service', false, false);
+```
+
+Grant the role whatever privileges your application needs.
+
+### 2. Provide your application with a Workload Identity
+
+The recommended way to provide a credential to a pod in AKS is [Microsoft Entra Workload ID](https://learn.microsoft.com/azure/aks/workload-identity-overview). It works by federating a Kubernetes service account with an Entra ID identity, so a JWT projected into the pod can be exchanged for an Entra access token without ever sharing a long-lived secret.
+
+The setup steps are (see [Deploy and configure an AKS cluster with Workload ID](https://learn.microsoft.com/azure/aks/workload-identity-deploy-cluster) for the full walkthrough):
+
+1. Enable the OIDC issuer and Workload Identity add-on on your AKS cluster.
+2. Create a user-assigned managed identity (or use an existing one).
+3. Create a [federated identity credential](https://learn.microsoft.com/entra/workload-id/workload-identity-federation) on the managed identity, with the cluster's OIDC issuer as the `issuer`, `api://AzureADTokenExchange` as the `audience`, and `system:serviceaccount:<namespace>:<service-account-name>` as the `subject`.
+4. Create a Kubernetes service account annotated with `azure.workload.identity/client-id: <managed-identity-client-id>` and use it for your application's pods.
+5. Add `azure.workload.identity/use: "true"` to your pod template's labels.
+
+Once that's in place, the Workload Identity mutating webhook automatically injects the following environment variables into your pod, which the gem uses to acquire access tokens:
+
+| Variable | Provided by |
+| --- | --- |
+| `AZURE_TENANT_ID` | Workload Identity webhook |
+| `AZURE_CLIENT_ID` | Workload Identity webhook (from the service-account annotation) |
+| `AZURE_FEDERATED_TOKEN_FILE` | Workload Identity webhook (path to the projected service-account JWT) |
+
+### 3. Configure your application to use the workload-identity token generator
+
+Set the non-standard `azure_workload_identity` connection parameter to a truthy value (e.g. `true`). When the gem sees this parameter on a connection, it acquires an Entra access token using the environment variables above and injects it as the database password before handing the connection string to `libpq`.
+
+You can set this parameter in
+
+* the query string of a connection URI:
+
+  ```
+  postgresql://my-app@my-server.postgres.database.azure.com:5432/appdb?sslmode=require&azure_workload_identity=true
+  ```
+
+* a `key=value` pair in a connection string:
+
+  ```
+  user=my-app host=my-server.postgres.database.azure.com port=5432 dbname=appdb sslmode=require azure_workload_identity=true
+  ```
+
+* a `key: value` pair in a connection hash:
+
+  ```ruby
+  PG.connect(
+    user: "my-app",
+    host: "my-server.postgres.database.azure.com",
+    port: 5432,
+    dbname: "appdb",
+    sslmode: "require",
+    azure_workload_identity: "true"
+  )
+  ```
+
+* `database.yml`, if you're using Rails:
+
+  ```yaml
+  production:
+    adapter: postgresql
+    username: my-app
+    host: my-server.postgres.database.azure.com
+    port: 5432
+    database: appdb
+    sslmode: require
+    azure_workload_identity: true
+  ```
+
+The gem caches access tokens in memory for their reported lifetime (with a small refresh threshold), so concurrent and back-to-back connections share a single token exchange. Tokens are refreshed automatically on demand before they expire.
+
+If the defaults don't match your environment (for example, you're using a sovereign cloud or a non-OSS-RDBMS scope), you can install a custom token generator at boot:
+
+```ruby
+PG::AzureWorkloadIdentity.auth_token_generator = PG::AzureWorkloadIdentity::AuthTokenGenerator.new(
+  identity_endpoint: "https://login.microsoftonline.us/<tenant-id>/oauth2/v2.0/token",
+  client_id: "<client-id>",
+  scope: "https://ossrdbms-aad.database.usgovcloudapi.net/.default",
+  grant_type: "client_credentials",
+  client_assertion_type: "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
+  federated_token_file: "/var/run/secrets/azure/tokens/azure-identity-token"
+)
+```
+
+The object you assign just needs to respond to `#call` and return a string token, so you can swap in an entirely different implementation when needed.
+
+### 4. Set `sslmode` to `verify-full` (recommended)
+
+Although Azure Database for PostgreSQL flexible server requires TLS by default (`sslmode=require`), the strongest setting is `sslmode=verify-full` — this also verifies the server's certificate chain and hostname, preventing man-in-the-middle attacks. You'll need to point `sslrootcert` at the Azure-issued certificate bundle; see [TLS certificates for Azure Database for PostgreSQL flexible server](https://learn.microsoft.com/azure/postgresql/flexible-server/concepts-networking-ssl-tls#downloading-root-ca-certificates-and-updating-application-clients-in-certificate-pinning-scenarios) for the current chain and download instructions. Note that this might no be possible when using private endpoints with entries in private DNS zones.
+
+## Development
+
+After checking out the repo, run `bundle install` to install dependencies.
+Then, run `bundle exec rake` to run the linter and unit tests.
+
+The acceptance suite spins up a local PostgreSQL container via Docker Compose and exercises the full connection flow (with WebMock stubbing the Entra token endpoint) — run it with:
+
+```console
+$ bundle exec rake test:acceptance
+```
+
+To release a new version:
+
+1. Update the version number in [version.rb](lib/pg/azure_workload_identity/version.rb), and run `bundle install` to update [Gemfile.lock](Gemfile.lock).
+2. Update [CHANGELOG.md](CHANGELOG.md).
+3. Submit the changes as a pull request.
+4. Once merged, tag the release and push the tag to GitHub.
+
+## Contributing
+
+Bug reports and pull requests are welcome [on GitHub](https://github.com/atpoint-cloud/pg-azure_workload_identity).
+
+## Credits
+
+This gem is built on the design and structure of [`pg-aws_rds_iam`](https://github.com/haines/pg-aws_rds_iam) which provides the equivalent functionality for AWS RDS IAM authentication.
+
+## License
+
+The gem is available as open source under the terms of the [MIT License](https://opensource.org/licenses/MIT).

data/lib/pg/azure_workload_identity/active_record_postgresql_adapter.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+require_relative "auth_token_injector"
+
+module PG
+  module AzureWorkloadIdentity
+    # Patch prepended to
+    # `ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.singleton_class`
+    # that wires Azure Workload Identity authentication into the
+    # `rails dbconsole` / `rails db` flow.
+    #
+    # The adapter's `dbconsole` class method is responsible for launching
+    # `psql` against the configured database; without this patch, `psql`
+    # would prompt for (or fail without) a password when the configuration
+    # opts into workload identity. The override generates an auth token and
+    # exposes it via `PGPASSWORD` before delegating to the original
+    # implementation.
+    module ActiveRecordPostgreSQLAdapter
+      # Sets `PGPASSWORD` on the environment to a freshly generated Azure
+      # Workload Identity auth token when the configuration enables it, then
+      # delegates to the original `dbconsole` implementation.
+      #
+      # @param config [ActiveRecord::DatabaseConfigurations::DatabaseConfig]
+      #   the configuration for the database `psql` is being launched
+      #   against.
+      # @return [void]
+      def dbconsole(config, options = {})
+        AuthTokenInjector.new.inject_into_psql_env! config.configuration_hash, ENV
+        super
+      end
+    end
+  end
+end

data/lib/pg/azure_workload_identity/active_record_postgresql_database_tasks.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+require_relative "auth_token_injector"
+
+module PG
+  module AzureWorkloadIdentity
+    # Patch prepended to
+    # `ActiveRecord::Tasks::PostgreSQLDatabaseTasks` that wires Azure
+    # Workload Identity authentication into rake tasks that shell out to
+    # `psql` (e.g. `db:structure:load`, `db:structure:dump`).
+    #
+    # ActiveRecord builds an environment hash for the `psql` subprocess via
+    # the private `#psql_env` method; this override generates an auth token
+    # and adds it to that hash as `PGPASSWORD` whenever the configuration
+    # opts into workload identity, so the `psql` invocation authenticates
+    # without prompting.
+    module ActiveRecordPostgreSQLDatabaseTasks
+      private
+
+      # Builds the environment hash for the `psql` subprocess, injecting a
+      # freshly generated Azure Workload Identity auth token as
+      # `PGPASSWORD` when the configuration enables it.
+      #
+      # @return [Hash{String => String}] the environment hash to pass to
+      #   the `psql` subprocess.
+      def psql_env
+        super.tap do |psql_env|
+          AuthTokenInjector.new.inject_into_psql_env! configuration_hash, psql_env
+        end
+      end
+    end
+  end
+end

data/lib/pg/azure_workload_identity/auth_token.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+require "json"
+
+require_relative "error"
+
+module PG
+  module AzureWorkloadIdentity
+    # Wraps a fetched OAuth access token together with the moment it was
+    # generated, and answers whether it is still safely usable. Validity
+    # is measured against the monotonic clock so the answer is not
+    # affected by wall-clock jumps (NTP slews, VM clock corrections).
+    class AuthToken
+      # Seconds before the reported expiry at which the token is
+      # considered stale, so callers proactively refresh before it
+      # actually expires in flight.
+      REFRESH_THRESHOLD_SECONDS = 60
+
+      # @return [String] the bearer access token.
+      attr_reader :access_token
+
+      # Parses a token response from the Azure AD token endpoint and builds
+      # an {AuthToken} instance.
+      #
+      # @param json [String] the raw JSON response body from the token   endpoint.
+      # @return [AuthToken] the parsed token.
+      # @raise [Error] when the JSON is invalid or lacks the expected fields.
+      def self.from_json(json)
+        JSON.parse(json).then do |data|
+          new(
+            access_token: data.fetch("access_token").to_s,
+            expires_in: data.fetch("expires_in").to_i
+          )
+        end
+      rescue JSON::ParserError => e
+        raise Error, "Failed to parse token response from JSON: #{e.message}"
+      rescue KeyError => e
+        raise Error, "Token response is missing key #{e.key} in #{e.receiver}"
+      end
+
+      # @param access_token [String] the bearer access token returned by
+      #   the token endpoint.
+      # @param expires_in [Integer, String] seconds-until-expiry as
+      #   reported by the token endpoint.
+      # @param refresh_threshold [Integer] seconds before the reported
+      #   expiry at which the token should be considered stale.
+      def initialize(
+        access_token:,
+        expires_in:,
+        refresh_threshold: REFRESH_THRESHOLD_SECONDS
+      )
+        @access_token = access_token
+        @expiry = expires_in
+        @generated_at = now
+        @refresh_threshold = refresh_threshold
+      end
+
+      # @return [Boolean] `true` if the token is still valid with the
+      #   refresh threshold applied.
+      def valid?
+        (now - @generated_at) < (@expiry - @refresh_threshold)
+      end
+
+      private
+
+      def now
+        Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      end
+    end
+  end
+end

data/lib/pg/azure_workload_identity/auth_token_generator.rb

@@ -0,0 +1,151 @@
+# frozen_string_literal: true
+
+require "net/http"
+require "openssl"
+require "uri"
+
+require_relative "auth_token"
+require_relative "error"
+
+module PG
+  module AzureWorkloadIdentity
+    # Exchanges a Kubernetes-projected federated identity token (as mounted
+    # into pods by AKS Workload Identity) for an Azure AD access token via the
+    # OAuth 2.0 client-credentials flow with a JWT bearer client assertion.
+    #
+    # Instances are callable via {#call}, which returns a cached access token
+    # while it is still valid (with a refresh threshold applied) and otherwise
+    # fetches a new one. Token fetches are guarded by a mutex so concurrent
+    # callers share a single in-flight request.
+    #
+    # Configuration may be passed explicitly to {#initialize} or populated
+    # from the standard Azure Workload Identity environment variables via
+    # {.default}.
+    class AuthTokenGenerator
+      IDENTITY_ENDPOINT = "https://login.microsoftonline.com/%<tenant_id>s/oauth2/v2.0/token"
+      SCOPE = "https://ossrdbms-aad.database.windows.net/.default"
+      GRANT_TYPE = "client_credentials"
+      CLIENT_ASSERTION_TYPE = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer"
+
+      # Builds an {AuthTokenGenerator} using the standard Azure Workload
+      # Identity environment variables (`AZURE_TENANT_ID`,
+      # `AZURE_CLIENT_ID`, `AZURE_FEDERATED_TOKEN_FILE`) and conventional
+      # OAuth defaults.
+      #
+      # @return [AuthTokenGenerator]
+      def self.default
+        new(
+          identity_endpoint: format(IDENTITY_ENDPOINT, tenant_id: ENV.fetch("AZURE_TENANT_ID")),
+          client_id: ENV.fetch("AZURE_CLIENT_ID"),
+          scope: SCOPE,
+          grant_type: GRANT_TYPE,
+          client_assertion_type: CLIENT_ASSERTION_TYPE,
+          federated_token_file: ENV.fetch("AZURE_FEDERATED_TOKEN_FILE")
+        )
+      end
+
+      # @return [String] the Azure AD token endpoint URL.
+      attr_reader :identity_endpoint
+
+      # @return [String] the AAD application/client id.
+      attr_reader :client_id
+
+      # @return [String] the requested OAuth scope.
+      attr_reader :scope
+
+      # @return [String] the OAuth grant type
+      #   (typically `"client_credentials"`).
+      attr_reader :grant_type
+
+      # @return [String] the client-assertion type
+      #   (typically `"urn:ietf:params:oauth:client-assertion-type:jwt-bearer"`).
+      attr_reader :client_assertion_type
+
+      # @return [String] absolute path to the file containing the
+      #   Kubernetes-projected federated identity JWT.
+      attr_reader :federated_token_file
+
+      # @param identity_endpoint [String] the Azure AD token endpoint.
+      # @param client_id [String] the AAD application id.
+      # @param scope [String] the requested OAuth scope.
+      # @param grant_type [String] the OAuth grant type.
+      # @param client_assertion_type [String] the client-assertion type.
+      # @param federated_token_file [String] path to the projected
+      #   federated identity JWT.
+      def initialize( # rubocop:disable Metrics/ParameterLists
+        identity_endpoint:,
+        client_id:,
+        scope:,
+        grant_type:,
+        client_assertion_type:,
+        federated_token_file:
+      )
+        @identity_endpoint = URI.parse(identity_endpoint)
+        @client_id = client_id
+        @scope = scope
+        @grant_type = grant_type
+        @client_assertion_type = client_assertion_type
+        @federated_token_file = federated_token_file
+        @mutex = Mutex.new
+        @token = nil
+      end
+
+      # Returns a currently valid Azure AD access token, fetching a new one
+      # from the identity endpoint when the cached token is missing or about
+      # to expire. Thread-safe: concurrent callers share a single in-flight
+      # fetch.
+      #
+      # @return [String] the bearer access token.
+      # @raise [Error] when the federated token cannot be read, the HTTP
+      #   request fails, the response is non-2xx, or the response body is
+      #   not valid JSON / lacks the expected fields.
+      def call
+        @mutex.synchronize do
+          @token = refresh unless @token&.valid?
+          @token.access_token
+        end
+      end
+
+      private
+
+      def refresh
+        response = request_token
+        return AuthToken.from_json(response.body) if response.is_a?(Net::HTTPSuccess)
+
+        raise Error, "Azure AD token endpoint responded with #{response.code} #{response.message}: #{response.body}"
+      end
+
+      def request_token
+        Net::HTTP.start(
+          @identity_endpoint.hostname.to_s,
+          @identity_endpoint.port,
+          use_ssl: @identity_endpoint.scheme == "https"
+        ) do |http|
+          http.request(build_request)
+        end
+      rescue Net::OpenTimeout, Net::ReadTimeout, Net::WriteTimeout,
+             SocketError, IOError, SystemCallError, OpenSSL::SSL::SSLError => e
+        raise Error, "Failed to reach Azure AD token endpoint: #{e.class}: #{e.message}"
+      end
+
+      def build_request
+        Net::HTTP::Post.new(@identity_endpoint).tap do |request|
+          request["Content-Type"] = "application/x-www-form-urlencoded"
+          request.body = URI.encode_www_form(
+            client_id: client_id,
+            scope: scope,
+            client_assertion_type: client_assertion_type,
+            client_assertion: read_federated_token,
+            grant_type: grant_type
+          )
+        end
+      end
+
+      def read_federated_token
+        File.read(federated_token_file).strip
+      rescue SystemCallError, IOError => e
+        raise Error, "Failed to read federated token file #{federated_token_file.inspect}: #{e.class}: #{e.message}"
+      end
+    end
+  end
+end

data/lib/pg/azure_workload_identity/auth_token_injector.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+require "pg"
+require_relative "connection_info"
+
+module PG
+  module AzureWorkloadIdentity
+    # AuthTokenInjector is responsible for injecting auth tokens into connection
+    # strings and PSQL environment variables when required.
+    class AuthTokenInjector
+      # Initializes a new AuthTokenInjector.
+      #
+      # The generator is resolved lazily — only on the code paths that
+      # actually need a token — so that constructing an injector is cheap
+      # and side-effect-free for connections that don't opt into workload
+      # identity.
+      #
+      # @param auth_token_generator [#call, nil] A callable that generates
+      #   an auth token. Defaults to the module-level generator at the time
+      #   a token is first needed.
+      def initialize(auth_token_generator: nil)
+        @auth_token_generator = auth_token_generator
+      end
+
+      # Injects an auth token into the given connection string as password if required.
+      #
+      # @param connection_string [String] The original connection string.
+      # @return [String] The modified connection string with the auth token injected if required.
+      def inject_into_connection_string(connection_string)
+        connection_info = ConnectionInfo.from_connection_string(connection_string)
+        return connection_string unless generate_auth_token?(connection_info)
+
+        connection_info.password = auth_token_generator.call
+        connection_info.to_s
+      end
+
+      # Injects an auth token into the given PSQL environment hash if required.
+      # This is used for the db:console rake task.
+      #
+      # @param configuration_hash [Hash] The ActiveRecord connection configuration hash from database.yml
+      # @param psql_env [Hash] The environment hash to be passed to the PSQL process.
+      def inject_into_psql_env!(configuration_hash, psql_env)
+        connection_info = ConnectionInfo.from_active_record_configuration_hash(configuration_hash)
+        return unless generate_auth_token?(connection_info)
+
+        psql_env["PGPASSWORD"] = auth_token_generator.call
+      end
+
+      private
+
+      # Resolves the auth token generator lazily, falling back to the
+      # module-level default only when a token is actually needed.
+      def auth_token_generator
+        @auth_token_generator ||= AzureWorkloadIdentity.auth_token_generator
+      end
+
+      # Determines whether an auth token should be generated based on the connection information.
+      #
+      # @param connection_info [ConnectionInfo] The parsed connection information.
+      # @return [Boolean] true if `azure_workload_identity` is set, false otherwise.
+      def generate_auth_token?(connection_info)
+        connection_info.azure_workload_identity
+      end
+    end
+  end
+end

data/lib/pg/azure_workload_identity/connection.rb

@@ -0,0 +1,83 @@
+# frozen_string_literal: true
+
+require_relative "auth_token_injector"
+require_relative "connection_info"
+
+module PG
+  module AzureWorkloadIdentity
+    # Patch prepended to `PG::Connection.singleton_class` that teaches libpq's
+    # connection-parameter machinery about the non-standard
+    # `azure_workload_identity` option and injects a freshly generated auth
+    # token as the password when that option is set.
+    #
+    # The three overrides cooperate:
+    # - {#conndefaults} adds `azure_workload_identity` to the list of known
+    #   parameters so it is recognized rather than rejected as unknown.
+    # - {#conninfo_parse} parses a connection string while preserving the
+    #   value of `azure_workload_identity` in the returned parameter list
+    #   (libpq would otherwise drop it).
+    # - {#parse_connect_args} runs the final connection string through
+    #   {AuthTokenInjector} so the password is replaced with a generated
+    #   token whenever `azure_workload_identity` is set.
+    module Connection
+      # Lists libpq's connection parameter defaults, with the
+      # `azure_workload_identity` option appended.
+      #
+      # @return [Array<Hash>] the connection parameter defaults.
+      def conndefaults
+        super + [conndefault_azure_workload_identity]
+      end
+
+      # Parses a connection string into the libpq parameter array, preserving
+      # the `azure_workload_identity` option.
+      #
+      # The input is first round-tripped through {ConnectionInfo} so the
+      # custom option is stripped before being handed to libpq's parser, and
+      # then re-appended to the parsed result with its original value.
+      #
+      # @param connection_string [String] either a URI-style or key=value
+      #   connection string.
+      # @return [Array<Hash>] the parsed parameter array, including the
+      #   `azure_workload_identity` entry when set.
+      def conninfo_parse(connection_string)
+        connection_info = ConnectionInfo.from_connection_string(connection_string)
+
+        super(connection_info.to_s).tap do |result|
+          if connection_info.azure_workload_identity
+            result << conndefault_azure_workload_identity.merge(val: connection_info.azure_workload_identity)
+          end
+        end
+      end
+
+      # Normalizes connection arguments and injects an Azure Workload Identity
+      # auth token as the password when `azure_workload_identity` is set.
+      #
+      # @param args [Array] the arguments passed to `PG::Connection.new` /
+      #   `PG.connect` (forwarded to `super`).
+      # @return [String] the connection string with the auth token injected
+      #   if applicable, otherwise unchanged.
+      def parse_connect_args(...)
+        AuthTokenInjector.new.inject_into_connection_string(super)
+      end
+
+      private
+
+      # Connection-parameter descriptor for the `azure_workload_identity`
+      # option, matching the shape of entries returned by libpq's
+      # `PQconndefaults`.
+      #
+      # @return [Hash] the descriptor hash.
+      def conndefault_azure_workload_identity
+        {
+          keyword: "azure_workload_identity",
+          envvar: nil,
+          compiled: nil,
+          val: nil,
+          label: "Azure-Workload-Identity",
+          dispchar: "",
+          dispsize: 64
+        }
+      end
+    end
+  end
+end

data/lib/pg/azure_workload_identity/connection_info/active_record_configuration_hash.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+module PG
+  module AzureWorkloadIdentity
+    module ConnectionInfo
+      # Wraps an ActiveRecord database configuration hash (as produced from
+      # `config/database.yml`) and exposes the connection parameters relevant
+      # to Azure Workload Identity authentication, including the non-standard
+      # `azure_workload_identity` flag.
+      #
+      # Accessor names follow this gem's connection-info contract (`user`,
+      # `host`, `port`) rather than ActiveRecord's keys (note the
+      # `:username` -> `#user` mapping).
+      class ActiveRecordConfigurationHash
+        # @param configuration_hash [Hash{Symbol => Object}] an ActiveRecord
+        #   database configuration hash (symbol keys, as found under an
+        #   environment in `database.yml`).
+        def initialize(configuration_hash)
+          @configuration_hash = configuration_hash
+        end
+
+        # @return [Object, nil] the value of the `azure_workload_identity` key,
+        #   or `nil` if not set. Truthy values opt the connection into Azure
+        #   Workload Identity authentication.
+        def azure_workload_identity
+          @configuration_hash[:azure_workload_identity]
+        end
+
+        # @return [String, nil] the database username (ActiveRecord's
+        #   `:username` key).
+        def user
+          @configuration_hash[:username]
+        end
+
+        # @return [String, nil] the database host.
+        def host
+          @configuration_hash[:host]
+        end
+
+        # @return [Integer, String, nil] the database port.
+        def port
+          @configuration_hash[:port]
+        end
+      end
+    end
+  end
+end

data/lib/pg/azure_workload_identity/connection_info/key_value_string.rb

@@ -0,0 +1,207 @@
+# frozen_string_literal: true
+
+require "strscan"
+require "pg"
+
+require_relative "../error"
+
+module PG
+  module AzureWorkloadIdentity
+    module ConnectionInfo
+      # Parses and manipulates a PostgreSQL connection string in libpq's
+      # key=value format (e.g. `"host=localhost user=admin port=5432"`).
+      #
+      # The parser preserves functional parity with libpq's `conninfo_parse`,
+      # supporting whitespace around the `=`, single-quoted values, and
+      # backslash escapes (`\\X` -> `X`) in both quoted and unquoted values.
+      #
+      # The non-standard `azure_workload_identity` parameter is recognized and
+      # extracted on construction so it can drive Azure Workload Identity
+      # authentication without being forwarded to libpq.
+      #
+      # @see https://github.com/postgres/postgres/blob/REL_18_4/src/interfaces/libpq/fe-connect.c#L6289-L6444
+      class KeyValueString
+        # @return [String, nil] the value of the `azure_workload_identity`
+        #   parameter that was present in the original connection string, or
+        #   `nil` if the parameter was not set.
+        attr_reader :azure_workload_identity
+
+        # Parses the given connection string and extracts the
+        # `azure_workload_identity` parameter (if any) so it is not forwarded to
+        # libpq.
+        #
+        # @param connection_string [String] a PostgreSQL key=value connection
+        #   string.
+        # @raise [Error] if the connection string is malformed.
+        def initialize(connection_string)
+          @params = Parser.new(connection_string).parse
+          @azure_workload_identity = @params.delete(:azure_workload_identity)
+        end
+
+        # @return [String, nil] the value of the `user` parameter.
+        def user
+          @params[:user]
+        end
+
+        # @return [String, nil] the value of the `host` parameter.
+        def host
+          @params[:host]
+        end
+
+        # @return [String, nil] the value of the `port` parameter.
+        def port
+          @params[:port]
+        end
+
+        # Sets the `password` parameter.
+        #
+        # @param value [String] the password to set.
+        # @return [String] the value that was set.
+        def password=(value)
+          @params[:password] = value
+        end
+
+        # Serializes the parameters back to a libpq key=value connection
+        # string. Values are quoted/escaped via `PG::Connection.quote_connstr`
+        # so they remain valid when whitespace or special characters are
+        # present. The `azure_workload_identity` parameter extracted at
+        # construction is not included.
+        #
+        # @return [String] the connection string.
+        def to_s
+          @params.map { |key, value| "#{key}=#{PG::Connection.quote_connstr(value)}" }.join(" ")
+        end
+
+        # @return [Hash{Symbol => String}] a copy of the parsed parameters
+        #   (excluding `azure_workload_identity`, which was extracted at
+        #   construction).
+        def to_h
+          @params.dup
+        end
+
+        # Tokenizes a libpq key=value connection string into a hash of
+        # symbol-keyed parameters.
+        #
+        # The grammar mirrors `conninfo_parse`: whitespace separates pairs,
+        # whitespace is allowed around `=`, values may be single-quoted, and
+        # `\\X` unescapes to `X` in both quoted and unquoted values. A
+        # trailing backslash at EOF inside an unquoted value is silently
+        # dropped (matching libpq).
+        #
+        # @see https://github.com/postgres/postgres/blob/REL_18_4/src/interfaces/libpq/fe-connect.c#L6289-L6444
+        class Parser
+          # @param connection_string [String] the connection string to parse.
+          def initialize(connection_string)
+            @buffer = StringScanner.new(connection_string)
+          end
+
+          # Parses the connection string passed to the constructor.
+          #
+          # @return [Hash{Symbol => String}] the parsed parameters, in the
+          #   order they appeared in the input. Later occurrences of the same
+          #   key overwrite earlier ones.
+          # @raise [Error] when a keyword is missing, `=` is missing
+          #   after a keyword, or a quoted value is not terminated.
+          def parse # rubocop:disable Metrics/MethodLength
+            result = {} # : Hash[Symbol, String]
+
+            skip_whitespace
+            until @buffer.eos?
+              key = parse_key
+              skip_whitespace
+              assert_followed_by_equals_sign key
+              skip_whitespace
+              value = parse_value
+              skip_whitespace
+
+              result[key] = value
+            end
+
+            @buffer.reset
+            result
+          end
+
+          private
+
+          # Advances the scanner past any run of whitespace characters.
+          #
+          # @return [Integer, nil] the number of characters consumed, or
+          #   `nil` if none matched.
+          def skip_whitespace
+            @buffer.skip(/\s+/)
+          end
+
+          # Reads the next keyword (any run of non-whitespace, non-`=`
+          # characters) and returns it as a symbol.
+          #
+          # @return [Symbol] the keyword.
+          # @raise [Error] if no keyword characters are present at the
+          #   current position.
+          def parse_key
+            key = @buffer.scan(/[^=\s]+/)
+            raise Error, %(Missing parameter name before "#{@buffer.rest}") unless key
+
+            key.to_sym
+          end
+
+          # Consumes a single `=` character at the current position.
+          #
+          # @param key [Symbol] the keyword that should be followed by `=`
+          #   (used in the error message).
+          # @raise [Error] if the next character is not `=`.
+          # @return [void]
+          def assert_followed_by_equals_sign(key)
+            raise Error, %(Missing "=" after "#{key}") unless @buffer.getch == "="
+          end
+
+          # Reads a value, dispatching to the quoted or unquoted form based
+          # on whether the next character is a single quote.
+          #
+          # @return [String] the parsed value (may be empty).
+          # @raise [Error] if a quoted value is not terminated.
+          def parse_value
+            if @buffer.peek(1) == "'"
+              parse_quoted_value
+            else
+              parse_unquoted_value
+            end
+          end
+
+          # Reads a single-quoted value, applying `\\X` -> `X` escape
+          # handling.
+          #
+          # @return [String] the value with quotes stripped and escapes
+          #   applied.
+          # @raise [Error] if the closing quote is missing.
+          def parse_quoted_value
+            value = +""
+            @buffer.skip(/'/)
+            loop do
+              value << (@buffer.scan(/[^'\\]+/) || "")
+              case @buffer.getch
+              when "'" then return value
+              when "\\" then value << (@buffer.getch || "")
+              else raise Error, "Unterminated quoted value"
+              end
+            end
+          end
+
+          # Reads an unquoted value, terminating at the next whitespace and
+          # applying `\\X` -> `X` escape handling. A trailing backslash at
+          # EOF is silently dropped (matching libpq).
+          #
+          # @return [String] the value with escapes applied (may be empty).
+          def parse_unquoted_value
+            value = +""
+            loop do
+              value << (@buffer.scan(/[^\s\\]+/) || "")
+              return value unless @buffer.getch == "\\"
+
+              value << (@buffer.getch || "")
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/pg/azure_workload_identity/connection_info/uri.rb

@@ -0,0 +1,91 @@
+# frozen_string_literal: true
+
+require "uri"
+
+module PG
+  module AzureWorkloadIdentity
+    module ConnectionInfo
+      # Parses and manipulates a PostgreSQL connection string in URI format
+      # (e.g. `postgres://user@host:5432/db?param=value`).
+      #
+      # Connection parameters may be provided either as URI components (userinfo,
+      # host, port) or as query-string parameters. Accessors prefer the URI
+      # component and fall back to the query string when absent.
+      #
+      # The non-standard `azure_workload_identity` query parameter is recognized
+      # and extracted on construction so the value can drive Azure Workload
+      # Identity authentication while keeping the rest of the connection string
+      # valid for libpq.
+      class URI
+        # Tests whether a connection string is in URI format.
+        #
+        # @param connection_string [String] the connection string to test.
+        # @return [Boolean] `true` if the string matches the RFC 2396 absolute
+        #   URI reference grammar, `false` otherwise.
+        def self.matches?(connection_string)
+          /\A#{::URI::RFC2396_PARSER.regexp[:ABS_URI_REF]}\z/.match?(connection_string)
+        end
+
+        # @return [String, nil] the value of the `azure_workload_identity` query
+        #   parameter that was present in the original connection string, or
+        #   `nil` if the parameter was not set.
+        attr_reader :azure_workload_identity
+
+        # Parses the given connection string and extracts the
+        # `azure_workload_identity` query parameter (if any) so it is not
+        # forwarded to libpq.
+        #
+        # @param connection_string [String] a PostgreSQL connection URI.
+        def initialize(connection_string)
+          @uri = ::URI.parse(connection_string)
+          @query = @uri.query ? ::URI.decode_www_form(@uri.query).to_h : {} # : Hash[String, String]
+          @azure_workload_identity = @query.delete("azure_workload_identity")
+        end
+
+        # @return [String, nil] the username from the URI userinfo, falling back
+        #   to the `user` query parameter.
+        def user
+          @uri.user || @query["user"]
+        end
+
+        # @return [String, nil] the host from the URI authority, falling back to
+        #   the `host` query parameter when the URI host is missing or empty.
+        def host
+          return @query["host"] if @uri.host.nil? || @uri.host.empty?
+
+          @uri.host
+        end
+
+        # @return [Integer, String, nil] the port from the URI authority,
+        #   falling back to the `port` query parameter.
+        def port
+          @uri.port || @query["port"]
+        end
+
+        # Sets the password, storing it as a `password` query parameter and
+        # clearing any password embedded in the URI userinfo.
+        #
+        # @param value [String] the password to set.
+        # @return [String] the value that was set.
+        def password=(value)
+          @uri.password = nil
+          @query["password"] = value
+        end
+
+        # Serializes the connection info back to a URI string.
+        #
+        # The query string is rebuilt from the current parameters (so the
+        # `azure_workload_identity` flag is omitted) and the scheme is normalized
+        # to always include the `"//"` authority separator, ensuring the result
+        # is a valid libpq connection URI.
+        #
+        # @return [String] the connection URI.
+        def to_s
+          @uri.query = ::URI.encode_www_form(@query)
+
+          @uri.to_s.sub(%r{^#{@uri.scheme}:(?!//)}, "#{@uri.scheme}://")
+        end
+      end
+    end
+  end
+end

data/lib/pg/azure_workload_identity/connection_info.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+require "pg"
+
+require_relative "connection_info/uri"
+require_relative "connection_info/key_value_string"
+require_relative "connection_info/active_record_configuration_hash"
+
+module PG
+  module AzureWorkloadIdentity
+    # Factory methods that wrap a PostgreSQL connection description (a
+    # connection string or an ActiveRecord configuration hash) in an object
+    # exposing a uniform accessor contract (`user`, `host`, `port`,
+    # `azure_workload_identity`, ...).
+    #
+    # Concrete return types are {URI}, {KeyValueString}, and
+    # {ActiveRecordConfigurationHash}.
+    module ConnectionInfo
+      # Wraps a PostgreSQL connection string, dispatching to the URI or
+      # key=value parser based on its format.
+      #
+      # @param connection_string [String] either a URI-style connection
+      #   string (e.g. `"postgres://user@host/db"`) or a libpq key=value
+      #   string (e.g. `"host=localhost user=admin"`).
+      # @return [URI, KeyValueString] a wrapper around the parsed
+      #   connection parameters.
+      def self.from_connection_string(connection_string)
+        if URI.matches?(connection_string)
+          URI.new(connection_string)
+        else
+          KeyValueString.new(connection_string)
+        end
+      end
+
+      # Wraps an ActiveRecord database configuration hash.
+      #
+      # @param configuration_hash [Hash{Symbol => Object}] an ActiveRecord
+      #   database configuration hash (as produced from `database.yml`).
+      # @return [ActiveRecordConfigurationHash] a wrapper around the
+      #   configuration hash.
+      def self.from_active_record_configuration_hash(configuration_hash)
+        ActiveRecordConfigurationHash.new(configuration_hash)
+      end
+    end
+  end
+end

data/lib/pg/azure_workload_identity/error.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+module PG
+  module AzureWorkloadIdentity
+    # Base error class for PG::AzureWorkloadIdentity-related errors.
+    class Error < StandardError
+    end
+  end
+end

data/lib/pg/azure_workload_identity/version.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+module PG
+  module AzureWorkloadIdentity
+    VERSION = "0.1.0"
+  end
+end

data/lib/pg/azure_workload_identity.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+require_relative "azure_workload_identity/auth_token_generator"
+require_relative "azure_workload_identity/connection"
+require_relative "azure_workload_identity/error"
+require_relative "azure_workload_identity/version"
+
+module PG
+  # Namespace for the pg-azure_workload_identity gem, which provides Azure
+  # Workload Identity authentication support for the pg gem and related tools.
+  module AzureWorkloadIdentity
+    # Returns the process-wide {AuthTokenGenerator}, lazily instantiated on
+    # first access via {AuthTokenGenerator.default} so that env-loading
+    # libraries (e.g. dotenv) have a chance to populate `ENV` before the
+    # generator captures its configuration.
+    #
+    # @return [AuthTokenGenerator]
+    def self.auth_token_generator
+      @auth_token_generator ||= AuthTokenGenerator.default
+    end
+
+    # Overrides the process-wide {AuthTokenGenerator}. Mainly useful for
+    # tests or for callers that want to provide a custom-configured
+    # generator instead of the env-driven default.
+    #
+    # @param generator [#call] any callable that returns an access token.
+    # @return [#call] the generator that was set.
+    def self.auth_token_generator=(generator)
+      @auth_token_generator = generator
+    end
+
+    PG::Connection.singleton_class.prepend Connection
+
+    if defined?(ActiveRecord)
+      require "active_record/connection_adapters/postgresql_adapter"
+
+      require_relative "azure_workload_identity/active_record_postgresql_adapter"
+      require_relative "azure_workload_identity/active_record_postgresql_database_tasks"
+
+      ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.singleton_class.prepend ActiveRecordPostgreSQLAdapter
+      ActiveRecord::Tasks::PostgreSQLDatabaseTasks.prepend ActiveRecordPostgreSQLDatabaseTasks
+    end
+  end
+end

data/pg-azure_workload_identity.gemspec

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+require_relative "lib/pg/azure_workload_identity/version"
+
+Gem::Specification.new do |spec|
+  spec.name = "pg-azure_workload_identity"
+  spec.version = PG::AzureWorkloadIdentity::VERSION
+  spec.authors = ["Simon Schmid"]
+  spec.email = ["simon@at-point.ch"]
+
+  spec.summary = "Workload identity authentication for Azure PostgreSQL flexible server/"
+  spec.homepage = "https://github.com/atpoint-cloud/pg-azure_workload_identity"
+  spec.license = "MIT"
+  spec.required_ruby_version = ">= 3.3"
+  spec.metadata["bug_tracker_uri"] = "#{spec.homepage}/issues"
+  spec.metadata["homepage_uri"] = spec.homepage
+  spec.metadata["source_code_uri"] = "#{spec.homepage}/blob/main"
+  spec.metadata["changelog_uri"] = "#{spec.homepage}/blob/main/CHANGELOG.md"
+
+  # Uncomment the line below to require MFA for gem pushes.
+  # This helps protect your gem from supply chain attacks by ensuring
+  # no one can publish a new version without multi-factor authentication.
+  # See: https://guides.rubygems.org/mfa-requirement-opt-in/
+  spec.metadata["rubygems_mfa_required"] = "true"
+
+  # Specify which files should be added to the gem when it is released.
+  spec.files = Dir[
+    "lib/**/*.rb",
+    ".yardopts",
+    "CHANGELOG.md",
+    "LICENSE.txt",
+    "README.md",
+    "pg-azure_workload_identity.gemspec"
+  ]
+
+  spec.require_paths = ["lib"]
+
+  # Uncomment to register a new dependency of your gem
+  spec.add_dependency "pg", "~> 1.5"
+
+  # For more information and examples about making a new gem, check out our
+  # guide at: https://guides.rubygems.org/make-your-own-gem/
+end

metadata

@@ -0,0 +1,76 @@
+--- !ruby/object:Gem::Specification
+name: pg-azure_workload_identity
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- Simon Schmid
+bindir: bin
+cert_chain: []
+date: 1980-01-02 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: pg
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.5'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.5'
+email:
+- simon@at-point.ch
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- ".yardopts"
+- CHANGELOG.md
+- LICENSE.txt
+- README.md
+- lib/pg/azure_workload_identity.rb
+- lib/pg/azure_workload_identity/active_record_postgresql_adapter.rb
+- lib/pg/azure_workload_identity/active_record_postgresql_database_tasks.rb
+- lib/pg/azure_workload_identity/auth_token.rb
+- lib/pg/azure_workload_identity/auth_token_generator.rb
+- lib/pg/azure_workload_identity/auth_token_injector.rb
+- lib/pg/azure_workload_identity/connection.rb
+- lib/pg/azure_workload_identity/connection_info.rb
+- lib/pg/azure_workload_identity/connection_info/active_record_configuration_hash.rb
+- lib/pg/azure_workload_identity/connection_info/key_value_string.rb
+- lib/pg/azure_workload_identity/connection_info/uri.rb
+- lib/pg/azure_workload_identity/error.rb
+- lib/pg/azure_workload_identity/version.rb
+- pg-azure_workload_identity.gemspec
+homepage: https://github.com/atpoint-cloud/pg-azure_workload_identity
+licenses:
+- MIT
+metadata:
+  bug_tracker_uri: https://github.com/atpoint-cloud/pg-azure_workload_identity/issues
+  homepage_uri: https://github.com/atpoint-cloud/pg-azure_workload_identity
+  source_code_uri: https://github.com/atpoint-cloud/pg-azure_workload_identity/blob/main
+  changelog_uri: https://github.com/atpoint-cloud/pg-azure_workload_identity/blob/main/CHANGELOG.md
+  rubygems_mfa_required: 'true'
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '3.3'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 4.0.10
+specification_version: 4
+summary: Workload identity authentication for Azure PostgreSQL flexible server/
+test_files: []