dwh 0.3.0 → 0.4.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: e70f914cc994c4be7a9d76b0d72d170ae6bf4d895427ec90adcdf9e3099774fe
-  data.tar.gz: 3ef66bc3d9a326bbae4b51d2bb3ec6af45425971617aee3a3131c7d496cf9127
+  metadata.gz: 9a67d957a140e258a8bef00efbae509bb7169e9de98dae911945316e7f24afb1
+  data.tar.gz: 0ae9acadd0a1a8e0e508f8a264c1534118d323202706c0488f3ae17ebdd98e9e
 SHA512:
-  metadata.gz: d05640e86dc5a6df2135173dd7a513df0805b4035fa5c4c5fa182659b1286fb6fd78f67ff2e259899f3d4ecab19b44bf5e7bfb427acfa5daa6065d71fb2fa18d
-  data.tar.gz: 28e5c623c8401dea1d222a1b318325543d4c209e083944d2fd43367ae6721c2ebe1c2b7e69190462a19887770d722edb7f2d1d851f90e5cc53b4c51cf2b65953
+  metadata.gz: 1aff93e7071cd35b748b43e174c3beeea2fb4c748a97bbf81d407e95e978c411e160b53ebe5cdef331ed685233407143485a0be7775b6b62db73538aebf05718
+  data.tar.gz: d9ad974f9e7a4edf0211b40bb0be33f730f8cd6abace42df05aa4339dcfc53babeea6c495c9d6d4d4b91519fa30f03852e99f6a0a0e14912914daf91bd62fb00

data/CHANGELOG.md

@@ -1,5 +1,35 @@
 ## [Unreleased]
 
+## [0.4.2] - 2026-05-22
+
+### Fixed
+
+- **DuckDB adapter**: Pass connection options with `config:` when opening a database so initialization uses the correct DuckDB API parameter.
+
+## [0.4.1] - 2026-04-29
+
+### Added
+
+- Databricks `execute_stream` support for `EXTERNAL_LINKS` result delivery using CSV downloads
+
+### Changed
+
+- Databricks now uses method-specific result delivery defaults: `execute` uses `INLINE` + `JSON_ARRAY`, and `execute_stream` uses `EXTERNAL_LINKS` + `CSV`.
+
+## [0.4.0] - 2026-04-28
+
+### Added
+
+- Token persistence interface via `DWH::TokenStore` for adapters that support OAuth token lifecycle management.
+- `TokenManageable` adapter concern for standardized token read/write behavior across adapters.
+- PKCE-based U2M OAuth support for the Databricks adapter.
+- Expanded tests for OAuth and Databricks token flows.
+
+### Changed
+
+- Databricks adapter now requires explicit `auth_mode` to reduce ambiguous auth configuration.
+- Updated documentation for adapter auth/token usage and adapter authoring.
+
 ## [0.3.0] - 2026-04-22
 
 ### Changed

data/docs/guides/adapters.md

@@ -179,7 +179,7 @@ adapter = DWH.create(:snowflake, {
   account_identifier: 'myorg-myaccount.us-east-1',
   oauth_client_id: '<YOUR_CLIENT_ID>',
   oauth_client_secret: '<YOUR_CLIENT_SECRET>',
-  oauth_redirect_url: 'https://localhost:3030/some/path',
+  oauth_redirect_uri: 'https://localhost:3030/some/path',
   database: 'ANALYTICS',
   client_name: 'myapp' # sent as user agent header value
 })
@@ -189,10 +189,119 @@ To successfully use OAuth you have to pass the adapter valid access and refresh
 
 The typical flow is like so:
 
-1. Generate an authorization code by visiting the url generated by `adapter.authorization_url.`  This will redirect to the configured `oauth_redirect_url.`  You must be able to retrieve the `code` from there.
+1. Generate an authorization code by visiting the url generated by `adapter.authorization_url.`  This will redirect to the configured `oauth_redirect_uri.`  You must be able to retrieve the `code` from there.
 2. Take the code from above and generate new access tokens: `adapter.generate_oauth_tokens(code)`. This will return Hash with access_token and refresh_token.  You can cache and reuse this until the refresh_token gets expired. This method will also apply the token to the current adapter instance.
 3. You can apply an existing set of tokens like so:`adapter.apply_oauth_tokens(access_token: token, refresh_token: token, expires_at: Time.now)`
 
+### Host integration contract (CLI / server)
+
+The host application is responsible for orchestration and persistence. DWH is responsible for OAuth protocol calls and token lifecycle methods.
+
+Public OAuth methods you can call on adapters that include `OpenAuthorizable`:
+
+- `authorization_url(state:, scope:)` - build provider authorize URL (authorization-code flow only)
+- `generate_oauth_tokens(code)` - exchange auth code for tokens and apply/store them
+- `apply_oauth_tokens(access_token:, refresh_token:, expires_at:)` - inject tokens from host storage
+- `oauth_access_token` - get a usable access token (load/refresh/mint as needed)
+- `refresh_access_token` - explicitly refresh using current refresh token
+- `mint_access_token` - explicitly mint via client credentials (M2M only)
+- `oauth_token_info` - inspect current token state
+
+Host requirements:
+
+1. Build and pass adapter config to `DWH.create(...)`, including correct `auth_mode` and OAuth fields.
+2. For U2M, implement browser redirect + callback capture, then call `generate_oauth_tokens(code)`.
+3. Validate OAuth `state` in the host callback handler (DWH does not enforce callback state verification).
+4. Persist tokens either by:
+   - passing a `token_store` object (`load`, `store`, `delete`), or
+   - storing tokens externally and rehydrating with `apply_oauth_tokens`.
+5. Handle auth exceptions and trigger reconnect UX when needed (for example, expired/invalid refresh token).
+
+Call order by mode:
+
+- **M2M (`oauth_m2m`)**
+  1. `adapter = DWH.create(...)`
+  2. Run query/test methods (`execute`, `test_connection`, etc.)
+  3. DWH internally calls `oauth_access_token`, which mints or refreshes as required
+
+- **U2M (`oauth_u2m`)**
+  1. `adapter = DWH.create(...)`
+  2. `url = adapter.authorization_url(...)`
+  3. Host sends user to URL and receives callback with `code`
+  4. `adapter.generate_oauth_tokens(code)`
+  5. Run query/test methods; DWH reuses and refreshes tokens as needed
+
+Note: for PKCE-enabled U2M providers, call `authorization_url` and `generate_oauth_tokens` on the same adapter instance.
+
+## Databricks
+
+The Databricks adapter uses the SQL Statements REST API and supports OAuth with
+both machine-to-machine (M2M) and user-to-machine (U2M) authorization-code flow.
+Set `auth_mode` explicitly to select the flow.
+
+### Basic configuration
+
+```ruby
+adapter = DWH.create(:databricks, {
+  host: 'workspace.cloud.databricks.com',
+  auth_mode: 'oauth_m2m',
+  warehouse: 'warehouse_id',
+  oauth_client_id: '<CLIENT_ID>',
+  oauth_client_secret: '<CLIENT_SECRET>',
+  catalog: 'main',
+  schema: 'default'
+})
+```
+
+### M2M (service principal) flow
+
+Set `auth_mode: 'oauth_m2m'`. The adapter mints access tokens using
+`grant_type=client_credentials`.
+
+### U2M (authorization code) flow
+
+Set `auth_mode: 'oauth_u2m'` and provide `oauth_redirect_uri` in config, then run:
+
+1. Generate authorize URL from `adapter.authorization_url`
+2. Capture `code` from redirect callback
+3. Exchange with `adapter.generate_oauth_tokens(code)`
+
+When U2M is active, PKCE is applied automatically by the adapter.
+
+### Large result sets (>25MiB)
+
+Databricks `INLINE` result disposition is limited for large payloads. The adapter uses
+method-driven defaults:
+
+- `execute` uses `INLINE` + `JSON_ARRAY` (in-memory/object style results)
+- `execute_stream` uses `EXTERNAL_LINKS` + `CSV` (large export path)
+
+```ruby
+adapter = DWH.create(:databricks, {
+  host: 'workspace.cloud.databricks.com',
+  auth_mode: 'oauth_m2m',
+  warehouse: 'warehouse_id',
+  oauth_client_id: '<CLIENT_ID>',
+  oauth_client_secret: '<CLIENT_SECRET>'
+})
+
+File.open('export.csv', 'w') do |io|
+  adapter.execute_stream('SELECT * FROM big_table', io)
+end
+```
+
+For low-memory exports, prefer `File`/`Tempfile`/pipes as the IO target. Avoid `StringIO`
+for very large result sets, because it keeps output bytes in memory.
+When using `execute_stream` with EXTERNAL_LINKS CSV, streaming stats/row counts are tracked
+consistently with other adapters (data rows only, header excluded).
+
+### Migration note
+
+Databricks now requires explicit `auth_mode`.
+
+- Existing service-principal setups should set `auth_mode: 'oauth_m2m'`
+- U2M setups should set `auth_mode: 'oauth_u2m'` and provide `oauth_redirect_uri`
+
 ## MySQL Adapter
 
 The MySQL adapter uses the `mysql2` gem. Note that MySQL's concept of "database" maps to "schema" in DWH.

data/docs/guides/creating-adapters.md

@@ -363,6 +363,49 @@ end
 
 ## Advanced Features
 
+### Identity-Bound Token Store Integration
+
+Adapters that use OAuth/M2M token exchange can support host-managed persistence by
+accepting `token_store` in adapter config and using base helpers from `Adapter`.
+
+`token_store` should be identity-bound by the host app (for example datasource-bound
+for service accounts, user+datasource-bound for per-user OAuth). The adapter should
+not parse or infer identity keys.
+
+```ruby
+class MyOAuthAdapter < Adapter
+  config :token_store, Object, required: false, default: nil
+
+  def access_token
+    payload = load_tokens_from_store
+    apply_token_payload(payload) if payload
+    return @access_token if @access_token && !token_expired?
+
+    token = request_new_token!
+    store_tokens_in_store(
+      access_token: token[:access_token],
+      refresh_token: token[:refresh_token],
+      expires_at: token[:expires_at]
+    )
+    @access_token
+  end
+end
+```
+
+Expected token-store methods:
+
+- `load` -> returns `nil` or hash with `access_token`, optional `refresh_token`, and `expires_at`
+- `store(token_hash)` -> persists latest token payload
+- `delete` -> optional revoke/cleanup hook for terminal auth failures
+
+For OAuth adapters, keep token persistence centralized and only override provider hooks:
+
+- `oauth_supports_authorization_code_flow?` -> `true` for U2M flows
+- `oauth_supports_client_credentials_flow?` -> `true` for M2M flows
+- `oauth_client_credentials_params` -> provider-specific client-credentials form body
+- `oauth_tokenization_url` -> provider token endpoint
+- `oauth_token_expiry_leeway_seconds` -> eager refresh buffer (for near-expiry tokens)
+
 ### Error Handling
 
 ```ruby

data/docs/guides/usage.md

@@ -341,6 +341,71 @@ readonly_analytics = DWH.create(:sqlite, {
 })
 ```
 
+## Identity-Bound Token Stores
+
+For OAuth and M2M adapters, DWH can optionally reuse tokens through a host-provided
+`token_store` object. The store should be identity-bound before it is passed into
+adapter config so DWH remains agnostic of app-level user/datasource models.
+
+### Token Store Contract
+
+`DWH::TokenStore` is the reference contract, but duck-typed objects are supported.
+
+```ruby
+class MyTokenStore < DWH::TokenStore
+  def load
+    # return nil or hash with:
+    # access_token (optional), refresh_token (optional), expires_at (optional)
+    #
+    # Notes:
+    # - providing expires_at enables proactive refresh/mint behavior
+    # - providing refresh_token enables refresh when access_token expires
+  end
+
+  def store(token)
+    # token includes at least access_token and expires_at
+  end
+
+  def delete
+    # optional cleanup/revoke path for terminal auth failures
+  end
+end
+```
+
+### Service Account (M2M) Example
+
+```ruby
+store = DatasourceTokenStore.new(datasource_id: datasource.id)
+
+client = DWH.create(:databricks, {
+  host: datasource.host,
+  auth_mode: 'oauth_m2m',
+  warehouse: datasource.warehouse,
+  oauth_client_id: datasource.oauth_client_id,
+  oauth_client_secret: datasource.oauth_client_secret,
+  token_store: store
+})
+```
+
+### Per-User OAuth Example (host-owned)
+
+```ruby
+store = UserDatasourceTokenStore.new(
+  user_id: current_user.id,
+  datasource_id: datasource.id
+)
+
+client = DWH.create(:snowflake, {
+  auth_mode: 'oauth',
+  account_identifier: datasource.account_identifier,
+  database: datasource.database,
+  oauth_client_id: datasource.oauth_client_id,
+  oauth_client_secret: datasource.oauth_client_secret,
+  oauth_redirect_uri: datasource.oauth_redirect_uri,
+  token_store: store
+})
+```
+
 ## Error Handling and Debugging
 
 ### Comprehensive Error Handling

data/lib/dwh/adapters/databricks.rb

@@ -1,11 +1,14 @@
 require 'csv'
-require 'base64'
+require_relative 'open_authorizable'
 
 module DWH
   module Adapters
     # Databricks adapter for executing SQL queries against Databricks SQL warehouses.
     #
-    # Supports OAuth M2M (service principal) authentication only.
+    # Supports OAuth M2M (service principal) and U2M (authorization code) flows.
+    # The host application must set auth_mode explicitly:
+    # - oauth_m2m: client_credentials flow
+    # - oauth_u2m: authorization_code + PKCE flow
     #
     # @example Connection with OAuth (service principal)
     #   DWH.create(:databricks, {
@@ -17,7 +20,15 @@ module DWH
     #     schema: 'default'
     #   })
     class Databricks < Adapter
+      include OpenAuthorizable
+
+      oauth_with authorize: ->(adapter) { "https://#{adapter.host}/oidc/v1/authorize" },
+                 tokenize: ->(adapter) { "https://#{adapter.host}/oidc/v1/token" },
+                 default_scope: 'all-apis'
+
       config :host, String, required: true, message: 'Databricks workspace host (e.g., adb-xxx.databricks.cloud.com)'
+      config :auth_mode, String, required: true, allowed: %w[oauth_m2m oauth_u2m],
+                                 message: 'Authentication mode: oauth_m2m or oauth_u2m'
       config :oauth_client_id, String, required: true, message: 'OAuth client ID (service principal application ID)'
       config :oauth_client_secret, String, required: true, message: 'OAuth client secret'
       config :client_name, String, required: false, default: 'Ruby DWH Gem', message: 'Client name sent to Databricks'
@@ -33,7 +44,7 @@ module DWH
 
       def initialize(config)
         super
-        validate_auth_config
+        validate_oauth_config
       end
 
       def connection
@@ -44,7 +55,7 @@ module DWH
           url: "https://#{workspace_host}",
           headers: {
             'Content-Type' => 'application/json',
-            'Authorization' => "Bearer #{auth_token}",
+            'Authorization' => "Bearer #{oauth_access_token}",
             'User-Agent' => config[:client_name]
           },
           request: {
@@ -67,7 +78,7 @@ module DWH
       def execute(sql, format: :array, retries: 0)
         result = with_retry(retries + 1) do
           with_debug(sql) do
-            response = submit_query(sql)
+            response = submit_query_for_execute(sql)
             fetch_data(handle_query_response(response))
           end
         end
@@ -78,7 +89,7 @@ module DWH
       def execute_stream(sql, io, stats: nil, retries: 0)
         with_retry(retries) do
           with_debug(sql) do
-            response = submit_query(sql)
+            response = submit_query_for_execute_stream(sql)
             fetch_data(handle_query_response(response), io: io, stats: stats)
           end
         end
@@ -92,7 +103,7 @@ module DWH
       # @yield [chunk] yields each chunk of data as it's processed
       def stream(sql, &block)
         with_debug(sql) do
-          response = submit_query(sql)
+          response = submit_query_for_execute(sql)
           fetch_data(handle_query_response(response), proc: block)
         end
       end
@@ -159,42 +170,14 @@ module DWH
 
       private
 
-      def validate_auth_config
-        raise ConfigError, 'oauth_client_id is required' unless config[:oauth_client_id]
-        raise ConfigError, 'oauth_client_secret is required' unless config[:oauth_client_secret]
-      end
-
-      def auth_token
-        return @oauth_access_token if @oauth_access_token && !token_expired?
-
-        request_oauth_access_token!
-        @oauth_access_token
-      end
-
-      def request_oauth_access_token!
-        credentials = Base64.strict_encode64("#{config[:oauth_client_id]}:#{config[:oauth_client_secret]}")
-        response = Faraday.post(
-          "https://#{workspace_host}/oidc/v1/token",
-          'grant_type=client_credentials&scope=all-apis',
-          'Authorization' => "Basic #{credentials}",
-          'Content-Type' => 'application/x-www-form-urlencoded'
-        )
-
-        raise AuthenticationError, "OAuth M2M token request failed (#{response.status}): #{response.body}" unless response.status == 200
-
-        data = JSON.parse(response.body)
-        @oauth_access_token = data['access_token']
-        expires_in = data['expires_in'] || 3600
-        @token_expires_at = Time.now + [expires_in - 60, 60].max
-      end
-
       def reset_connection
         @oauth_access_token = nil
+        @oauth_refresh_token = nil
         @token_expires_at = nil
         close
       end
 
-      def submit_query(sql)
+      def submit_query(sql, disposition: 'INLINE', format: 'JSON_ARRAY')
         connection.post(STATEMENTS_API) do |req|
           req.body = {
             statement: sql,
@@ -203,12 +186,20 @@ module DWH
             schema: config[:schema],
             wait_timeout: '30s',
             on_wait_timeout: 'CONTINUE',
-            format: 'JSON_ARRAY',
-            disposition: 'INLINE'
+            format:,
+            disposition:
           }.compact.merge(extra_query_params).to_json
         end
       end
 
+      def submit_query_for_execute(sql)
+        submit_query(sql, disposition: 'INLINE', format: 'JSON_ARRAY')
+      end
+
+      def submit_query_for_execute_stream(sql)
+        submit_query(sql, disposition: 'EXTERNAL_LINKS', format: 'CSV')
+      end
+
       def handle_query_response(response)
         body = JSON.parse(response.body)
 
@@ -250,6 +241,8 @@ module DWH
       end
 
       def fetch_data(result, io: nil, stats: nil, proc: nil)
+        return fetch_external_links_data(result, io:, stats:, proc:) if result.dig('result', 'external_links')
+
         columns = result.dig('manifest', 'schema', 'columns')&.map { |col| col['name'] } || []
         chunks = result.dig('manifest', 'chunks') || []
         collector = {
@@ -279,6 +272,60 @@ module DWH
         collector
       end
 
+      def fetch_external_links_data(result, io:, proc:, stats: nil)
+        if io.nil?
+          raise UnsupportedCapability,
+                "Databricks EXTERNAL_LINKS is supported only for execute_stream. Use result_format: 'CSV' with execute_stream."
+        end
+        raise UnsupportedCapability, 'Databricks EXTERNAL_LINKS does not support stream/yield. Use execute_stream.' if proc
+
+        csv_buffer = +''
+        header_skipped = false
+        current = result
+        loop do
+          links = current.dig('result', 'external_links') || []
+          links.each do |link|
+            url = link['external_link']
+            raise ExecutionError, 'Databricks external link missing external_link URL' if url.to_s.strip.empty?
+
+            response = external_link_http_client.get(url)
+            raise ExecutionError, "Failed to download Databricks external link: #{response.status}" unless response.status == 200
+
+            body = response.body.to_s
+            io << body
+            header_skipped = append_csv_stats(stats, body, csv_buffer, header_skipped:) if stats
+          end
+
+          next_chunk_internal_link = links.first&.dig('next_chunk_internal_link')
+          break if next_chunk_internal_link.to_s.strip.empty?
+
+          current = JSON.parse(connection.get(next_chunk_internal_link).body)
+        end
+
+        io.rewind
+        { columns: [], data: [], io: io, stats: stats, wrote_header: true }
+      end
+
+      def append_csv_stats(stats, chunk, csv_buffer, header_skipped:)
+        return if stats.nil?
+
+        csv_buffer << chunk
+        rows = CSV.parse(csv_buffer, skip_blanks: true)
+        rows.each_with_index do |row, index|
+          if !header_skipped && index.zero?
+            header_skipped = true
+            next
+          end
+
+          stats << row
+        end
+        csv_buffer.clear
+        header_skipped
+      rescue CSV::MalformedCSVError
+        logger.debug("Unparseable:\n #{chunk}")
+        header_skipped
+      end
+
       def write_data(data, collector, io = nil, stats = nil, proc = nil)
         if io
           unless collector[:wrote_header]
@@ -321,7 +368,38 @@ module DWH
       end
 
       def workspace_host
-        config[:host].to_s.gsub(%r{\Ahttps?://}, '').gsub(%r{/+\z}, '')
+        config[:host].to_s
+      end
+
+      def external_link_http_client
+        @external_link_http_client ||= Faraday.new
+      end
+
+      def oauth_supports_authorization_code_flow?
+        auth_mode == 'oauth_u2m'
+      end
+
+      def oauth_supports_client_credentials_flow?
+        auth_mode == 'oauth_m2m'
+      end
+
+      def oauth_redirect_uri_required?
+        oauth_supports_authorization_code_flow?
+      end
+
+      def oauth_client_credentials_params
+        {
+          grant_type: 'client_credentials',
+          scope: 'all-apis'
+        }
+      end
+
+      def oauth_token_expiry_leeway_seconds
+        30
+      end
+
+      def oauth_uses_pkce?
+        oauth_supports_authorization_code_flow?
       end
     end
   end

data/lib/dwh/adapters/duck_db.rb

@@ -32,7 +32,7 @@ module DWH
               ducked_config[key.to_s] = val
             end
           end
-          @db = DuckDB::Database.open(config[:file], ducked_config)
+          @db = DuckDB::Database.open(config[:file], config: ducked_config)
           self.class.databases[config[:file]] = @db
         end
 

data/lib/dwh/adapters/open_authorizable.rb

@@ -1,5 +1,7 @@
 require 'base64'
 require 'securerandom'
+require 'digest'
+require_relative 'token_manageable'
 
 module DWH
   module Adapters
@@ -26,7 +28,7 @@ module DWH
     module OpenAuthorizable
       # rubcop:disable Style/DocumentationModule
       module ClassMethods
-        def oauth_with(authorize:, tokenize:, default_scope: 'refresh_token')
+        def oauth_with(authorize: nil, tokenize: nil, default_scope: 'refresh_token')
           @oauth_settings = { authorize: authorize, tokenize: tokenize, default_scope: default_scope }
         end
 
@@ -39,14 +41,18 @@ module DWH
 
       def self.included(base)
         base.extend(ClassMethods)
+        base.include(TokenManageable)
         base.config :oauth_client_id, String, required: false, message: 'OAuth client_id'
         base.config :oauth_client_secret, String, required: false, message: 'OAuth client_secret'
         base.config :oauth_redirect_uri, String, required: false, message: 'OAuth redirect_uri'
-        base.config :oauth_scope, String, required: false, message: 'OAuth redirect_url'
+        base.config :oauth_scope, String, required: false, message: 'OAuth scope'
       end
 
       # Generate authorization URL for user to visit
       def authorization_url(state: SecureRandom.hex(16), scope: nil)
+        raise UnsupportedCapability, "#{adapter_name} does not support authorization-code OAuth flow" unless oauth_supports_authorization_code_flow?
+
+        code_verifier = oauth_pkce_code_verifier_for_session
         params = {
           'response_type' => 'code',
           'client_id' => oauth_client_id,
@@ -54,6 +60,7 @@ module DWH
           'state' => state,
           'scope' => scope || oauth_scope || oauth_settings[:default_scope]
         }.compact
+        params.merge!(oauth_pkce_authorization_params(code_verifier))
 
         uri = URI(oauth_settings[:authorize])
         uri.query = URI.encode_www_form(params)
@@ -66,7 +73,7 @@ module DWH
       #
       # param access_token [String] the access token
       # @param refresh_token [String] optional refresh token
-      def apply_oauth_tokens(access_token:, refresh_token: nil, expires_at: nil)
+      def apply_oauth_tokens(access_token: nil, refresh_token: nil, expires_at: nil)
         @oauth_access_token = access_token
         @oauth_refresh_token = refresh_token
         @token_expires_at = expires_at
@@ -77,11 +84,15 @@ module DWH
       # @param authorization_code [String] this code should come from
       #   the redirect that is captured from the #authorization_url
       def generate_oauth_tokens(authorization_code)
+        raise UnsupportedCapability, "#{adapter_name} does not support authorization-code OAuth flow" unless oauth_supports_authorization_code_flow?
+
+        code_verifier = oauth_pkce_code_verifier_for_session
         params = {
           grant_type: 'authorization_code',
           code: authorization_code,
           redirect_uri: oauth_redirect_uri
         }
+        params.merge!(oauth_pkce_token_params(code_verifier))
 
         response = oauth_http_client.post(oauth_tokenization_url) do |req|
           req.headers['Content-Type'] = 'application/x-www-form-urlencoded'
@@ -95,14 +106,23 @@ module DWH
       def refresh_access_token
         raise AuthenticationError, 'No refresh token available' unless @oauth_refresh_token
 
-        params = {
-          grant_type: 'refresh_token',
-          refresh_token: @oauth_refresh_token
-        }
+        params = oauth_refresh_token_params
+        response = oauth_http_client.post(oauth_tokenization_url) do |req|
+          req.headers['Content-Type'] = 'application/x-www-form-urlencoded'
+          req.headers['Authorization'] = oauth_token_request_auth_header
+          req.body = URI.encode_www_form(params)
+        end
 
+        oauth_token_response(response)
+      end
+
+      def mint_access_token
+        raise UnsupportedCapability, "#{adapter_name} does not support client-credentials OAuth flow" unless oauth_supports_client_credentials_flow?
+
+        params = oauth_client_credentials_params
         response = oauth_http_client.post(oauth_tokenization_url) do |req|
           req.headers['Content-Type'] = 'application/x-www-form-urlencoded'
-          req.headers['Authorization'] = basic_auth_header
+          req.headers['Authorization'] = oauth_token_request_auth_header
           req.body = URI.encode_www_form(params)
         end
 
@@ -116,22 +136,22 @@ module DWH
       # @return [String] access token
       # @raise [AuthenticationError]
       def oauth_access_token
-        if token_expired? && @oauth_refresh_token
-          refresh_access_token
+        load_oauth_tokens_from_store! unless @oauth_access_token || @oauth_refresh_token
+        return @oauth_access_token if oauth_token_usable?
 
-          # return token unless exception was raised
-          @oauth_access_token
-        elsif @oauth_access_token
-          @oauth_access_token
-        else
-          raise AuthenticationError,
-                'Access token was never set. Either run the auth flow or set the tokens via apply_oauth_tokens method.'
-        end
+        refresh_access_token if oauth_refresh_token_usable?
+        return @oauth_access_token if oauth_token_usable?
+
+        mint_access_token if oauth_supports_client_credentials_flow?
+        return @oauth_access_token if oauth_token_usable?
+
+        raise AuthenticationError,
+              'Access token was never set. Either run the auth flow, mint via client credentials, or set tokens via apply_oauth_tokens.'
       end
 
       # Check if we have a valid access token
       def oauth_authenticated?
-        @oauth_access_token && !token_expired?
+        @oauth_access_token && oauth_token_usable?
       end
 
       # Get current state of tokens
@@ -140,7 +160,7 @@ module DWH
           access_token: @oauth_access_token,
           refresh_token: @oauth_refresh_token,
           expires_at: @token_expires_at,
-          expired: token_expired?,
+          expired: !oauth_token_usable?,
           authenticated: oauth_authenticated?
         }
       end
@@ -148,9 +168,11 @@ module DWH
       def validate_oauth_config
         raise ConfigError, 'Missing config: oauth_client_id. Required for OAuth.' unless config[:oauth_client_id]
         raise ConfigError, 'Missing config: oauth_client_secret. Required for OAuth.' unless config[:oauth_client_secret]
-        raise ConfigError, 'Missing config: oauth_redirect_url. Required for OAuth.' unless config[:oauth_redirect_uri]
 
-        oauth_settings
+        raise ConfigError, 'Missing config: oauth_redirect_uri. Required for OAuth.' if oauth_redirect_uri_required? && !config[:oauth_redirect_uri]
+
+        oauth_settings if oauth_supports_authorization_code_flow?
+        true
       end
 
       def oauth_settings
@@ -170,6 +192,75 @@ module DWH
         "Basic #{credentials}"
       end
 
+      def oauth_token_request_auth_header
+        basic_auth_header
+      end
+
+      def oauth_refresh_token_params
+        {
+          grant_type: 'refresh_token',
+          refresh_token: @oauth_refresh_token
+        }
+      end
+
+      def oauth_client_credentials_params
+        {
+          grant_type: 'client_credentials',
+          scope: oauth_scope || oauth_settings[:default_scope]
+        }.compact
+      end
+
+      def oauth_supports_authorization_code_flow?
+        true
+      end
+
+      def oauth_supports_client_credentials_flow?
+        false
+      end
+
+      def oauth_redirect_uri_required?
+        oauth_supports_authorization_code_flow?
+      end
+
+      def oauth_token_expiry_leeway_seconds
+        0
+      end
+
+      # PKCE is optional and disabled by default
+      def oauth_uses_pkce?
+        false
+      end
+
+      def oauth_pkce_code_challenge_method
+        'S256'
+      end
+
+      def oauth_pkce_code_verifier_for_session
+        return nil unless oauth_uses_pkce?
+
+        @oauth_pkce_code_verifier_for_session ||= oauth_pkce_code_verifier
+      end
+
+      def oauth_pkce_code_verifier
+        SecureRandom.urlsafe_base64(64).delete('=')
+      end
+
+      def oauth_token_usable?
+        return false unless @oauth_access_token
+
+        !token_expiring_soon?
+      end
+
+      def oauth_refresh_token_usable?
+        @oauth_refresh_token && token_expired?
+      end
+
+      def token_expiring_soon?(seconds = oauth_token_expiry_leeway_seconds)
+        return true if @token_expires_at.nil?
+
+        (Time.now + seconds) >= @token_expires_at
+      end
+
       def oauth_http_client
         @oauth_http_client ||= Faraday.new(
           headers: {
@@ -191,11 +282,20 @@ module DWH
           # Calculate expiration time
           expires_in = data['expires_in'] || 3600
           @token_expires_at = Time.now + expires_in
+          store_tokens_in_store(
+            access_token: @oauth_access_token,
+            refresh_token: @oauth_refresh_token,
+            expires_at: @token_expires_at,
+            token_type: data['token_type'],
+            scope: data['scope'],
+            raw: data
+          )
 
           { success: true, data: data }
         else
           error_data = parse_error_response(response)
           if error_data['error'] == 'invalid_grant' && @oauth_refresh_token
+            delete_tokens_from_store
             raise TokenExpiredError, "Potentially expired refresh token. #{error_data['message']}"
           end
 
@@ -205,11 +305,48 @@ module DWH
 
       private
 
+      def load_oauth_tokens_from_store!
+        payload = load_tokens_from_store
+        return unless payload
+
+        apply_oauth_tokens(
+          access_token: payload[:access_token],
+          refresh_token: payload[:refresh_token],
+          expires_at: payload[:expires_at]
+        )
+      end
+
       def parse_error_response(response)
         JSON.parse(response.body)
       rescue JSON::ParserError
         { 'error' => 'unknown', 'message' => response.body }
       end
+
+      def oauth_pkce_authorization_params(code_verifier)
+        return {} unless oauth_uses_pkce?
+
+        {
+          'code_challenge' => oauth_pkce_code_challenge(code_verifier),
+          'code_challenge_method' => oauth_pkce_code_challenge_method
+        }
+      end
+
+      def oauth_pkce_token_params(code_verifier)
+        return {} unless oauth_uses_pkce?
+
+        { code_verifier: code_verifier }
+      end
+
+      def oauth_pkce_code_challenge(code_verifier)
+        case oauth_pkce_code_challenge_method
+        when 'S256'
+          Base64.urlsafe_encode64(Digest::SHA256.digest(code_verifier), padding: false)
+        when 'plain'
+          code_verifier
+        else
+          raise ConfigError, "Unsupported PKCE code challenge method: #{oauth_pkce_code_challenge_method}"
+        end
+      end
     end
   end
 end

data/lib/dwh/adapters/snowflake.rb

@@ -39,7 +39,7 @@ module DWH
     #     account_identifier: 'myorg-myaccount.us-east-1',
     #     oauth_client_id: '<YOUR_CLIENT_ID>',
     #     oauth_client_secret: '<YOUR_CLIENT_SECRET>',
-    #     oauth_redirect_url: 'https://localhost:3030/some/path',
+    #     oauth_redirect_uri: 'https://localhost:3030/some/path',
     #     database: 'ANALYTICS'
     #   })
     #

data/lib/dwh/adapters/token_manageable.rb

@@ -0,0 +1,81 @@
+require 'time'
+
+module DWH
+  module Adapters
+    # TokenManageable hold the logic to load, store and delete tokens from the token store.
+    module TokenManageable
+      def token_store
+        config[:token_store]
+      end
+
+      def load_tokens_from_store
+        return nil unless token_store.respond_to?(:load)
+
+        payload = token_store.load
+        normalize_token_payload(payload)
+      rescue StandardError => e
+        logger.warn("Failed loading token from token_store: #{e.message}")
+        nil
+      end
+
+      def store_tokens_in_store(token_payload)
+        return unless token_store.respond_to?(:store)
+
+        token_store.store(normalize_token_payload_for_store(token_payload))
+      rescue StandardError => e
+        logger.warn("Failed storing token in token_store: #{e.message}")
+      end
+
+      def delete_tokens_from_store
+        return unless token_store.respond_to?(:delete)
+
+        token_store.delete
+      rescue StandardError => e
+        logger.warn("Failed deleting token from token_store: #{e.message}")
+      end
+
+      private
+
+      def normalize_token_payload(payload)
+        return nil unless payload.is_a?(Hash)
+
+        data = payload.transform_keys(&:to_sym)
+        access_token = data[:access_token]
+        access_token = nil if access_token.to_s.strip == ''
+
+        refresh_token = data[:refresh_token]
+        refresh_token = nil if refresh_token.respond_to?(:empty?) && refresh_token.empty?
+        return nil if access_token.nil? && refresh_token.nil?
+
+        {
+          access_token: access_token&.to_s,
+          refresh_token: refresh_token,
+          expires_at: parse_token_expiry(data[:expires_at])
+        }
+      end
+
+      def normalize_token_payload_for_store(payload)
+        data = payload.transform_keys(&:to_sym)
+        cleaned = {
+          access_token: data[:access_token]&.to_s,
+          refresh_token: data[:refresh_token],
+          expires_at: parse_token_expiry(data[:expires_at]),
+          token_type: data[:token_type],
+          scope: data[:scope],
+          issued_at: parse_token_expiry(data[:issued_at]),
+          raw: data[:raw]
+        }
+        cleaned.reject { |_k, v| v.nil? }
+      end
+
+      def parse_token_expiry(value)
+        return value if value.is_a?(Time)
+        return nil if value.nil? || value.to_s.strip == ''
+
+        Time.parse(value.to_s)
+      rescue StandardError
+        nil
+      end
+    end
+  end
+end

data/lib/dwh/adapters.rb

@@ -79,6 +79,13 @@ module DWH
       # @return [Hash] the actual instance configuration
       attr_reader :config
 
+      # Optional host-implemented token store for OAuth token reuse.
+      # That should implement the following methods:
+      #   - load -> Hash|nil
+      #   - store(token_hash)
+      #   - delete
+      config :token_store, Object, required: false, default: nil, message: 'Token store instance implementing load/store/delete'
+
       def initialize(config)
         @config = config.transform_keys(&:to_sym)
         # Per instance customization of general settings

data/lib/dwh/token_store.rb

@@ -0,0 +1,24 @@
+module DWH
+  # Optional contract for host applications that want token persistence.
+  #
+  # The store instance should be identity-bound before it is passed into
+  # adapter config so the adapter remains unaware of user/datasource identity.
+  #
+  # This class is intentionally minimal and can be subclassed or duck-typed.
+  class TokenStore
+    # @return [Hash,nil] token payload
+    def load
+      raise NotImplementedError, "#{self.class} must implement ##{__method__}"
+    end
+
+    # @param token [Hash] normalized payload with at least access_token and expires_at
+    def store(_token)
+      raise NotImplementedError, "#{self.class} must implement ##{__method__}"
+    end
+
+    # Remove/revoke persisted token state.
+    def delete
+      raise NotImplementedError, "#{self.class} must implement ##{__method__}"
+    end
+  end
+end

data/lib/dwh/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module DWH
-  VERSION = '0.3.0'
+  VERSION = '0.4.2'
 end

data/lib/dwh.rb

@@ -5,6 +5,7 @@ require_relative 'dwh/errors'
 require_relative 'dwh/logger'
 require_relative 'dwh/streaming_stats'
 require_relative 'dwh/factory'
+require_relative 'dwh/token_store'
 require_relative 'dwh/adapters'
 require_relative 'dwh/table'
 require_relative 'dwh/table_stats'

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: dwh
 version: !ruby/object:Gem::Version
-  version: 0.3.0
+  version: 0.4.2
 platform: ruby
 authors:
 - Ajo Abraham
@@ -164,6 +164,7 @@ files:
 - lib/dwh/adapters/snowflake.rb
 - lib/dwh/adapters/sql_server.rb
 - lib/dwh/adapters/sqlite.rb
+- lib/dwh/adapters/token_manageable.rb
 - lib/dwh/adapters/trino.rb
 - lib/dwh/behaviors.rb
 - lib/dwh/capabilities.rb
@@ -192,6 +193,7 @@ files:
 - lib/dwh/streaming_stats.rb
 - lib/dwh/table.rb
 - lib/dwh/table_stats.rb
+- lib/dwh/token_store.rb
 - lib/dwh/version.rb
 - sig/dwh.rbs
 homepage: https://www.strata.site