Apply 3 suggestion(s) to 2 file(s)

Co-authored-by: Ravi Kumar <rkumar@gitlab.com>

Gemfile.lock

@@ -27,6 +27,9 @@ PATH
   remote: gems/gitlab-active-context
   specs:
     gitlab-active-context (0.0.1)
+      activesupport
+      connection_pool
+      pg
       zeitwerk
 
 PATH

Gemfile.next.lock

@@ -27,6 +27,9 @@ PATH
   remote: gems/gitlab-active-context
   specs:
     gitlab-active-context (0.0.1)
+      activesupport
+      connection_pool
+      pg
       zeitwerk
 
 PATH

gems/gitlab-active-context/Gemfile.lock

@@ -2,6 +2,9 @@ PATH
   remote: .
   specs:
     gitlab-active-context (0.0.1)
+      activesupport
+      connection_pool
+      pg
       zeitwerk
 
 GEM
@@ -87,6 +90,7 @@ GEM
     parser (3.3.6.0)
       ast (~> 2.4.1)
       racc
+    pg (1.5.9)
     psych (5.2.1)
       date
       stringio

gems/gitlab-active-context/gitlab-active-context.gemspec

@@ -19,6 +19,9 @@ Gem::Specification.new do |spec|
   spec.files = Dir['lib/**/*.rb']
   spec.require_paths = ["lib"]
 
+  spec.add_dependency 'activesupport'
+  spec.add_dependency 'connection_pool'
+  spec.add_dependency 'pg'
   spec.add_dependency 'zeitwerk'
 
   spec.add_development_dependency 'gitlab-styles'

gems/gitlab-active-context/lib/active_context.rb

 # frozen_string_literal: true
 
-require "zeitwerk"
+require 'active_support/core_ext/string/inflections'
+require 'active_support/core_ext/module/delegation'
+require 'connection_pool'
+require 'pg'
+require 'zeitwerk'
 loader = Zeitwerk::Loader.for_gem
 loader.setup
 
@@ -8,4 +12,12 @@ module ActiveContext
   def self.configure(...)
     ActiveContext::Config.configure(...)
   end
+
+  def self.config
+    ActiveContext::Config.current
+  end
+
+  def self.adapter
+    ActiveContext::Adapter.current
+  end
 end

gems/gitlab-active-context/lib/active_context/adapter.rb

+# frozen_string_literal: true
+
+module ActiveContext
+  module Adapter
+    class << self
+      def current
+        @current ||= load_adapter
+      end
+
+      private
+
+      def load_adapter
+        config = ActiveContext::Config.current
+        return nil unless config.enabled
+
+        name, hash = config.databases.first
+        return nil unless name
+
+        adapter = hash.fetch(:adapter)
+        return nil unless adapter
+
+        adapter_klass = adapter.safe_constantize
+        return nil unless adapter_klass
+
+        options = hash.fetch(:options)
+
+        adapter_klass.new(options)
+      end
+    end
+  end
+end

gems/gitlab-active-context/lib/active_context/config.rb

@@ -2,27 +2,27 @@
 
 module ActiveContext
   class Config
-    CONFIG = Struct.new(:enabled, :databases, :logger)
+    Cfg = Struct.new(:enabled, :databases, :logger)
 
     class << self
       def configure(&block)
         @instance = new(block)
       end
 
-      def config
-        @instance&.config || {}
+      def current
+        @instance&.config || Cfg.new
       end
 
       def enabled?
-        config.enabled || false
+        current.enabled || false
       end
 
       def databases
-        config.databases || {}
+        current.databases || {}
       end
 
       def logger
-        config.logger || Logger.new($stdout)
+        current.logger || Logger.new($stdout)
       end
     end
 
@@ -31,7 +31,7 @@ def initialize(config_block)
     end
 
     def config
-      struct = CONFIG.new
+      struct = Cfg.new
       @config_block.call(struct)
       struct
     end

gems/gitlab-active-context/lib/active_context/databases/concerns/adapter.rb

+# frozen_string_literal: true
+
+module ActiveContext
+  module Databases
+    module Concerns
+      module Adapter
+        attr_reader :client
+
+        delegate :search, to: :client
+
+        def initialize(options)
+          @client = client_klass.new(options)
+        end
+
+        def client_klass
+          raise NotImplementedError
+        end
+      end
+    end
+  end
+end

gems/gitlab-active-context/lib/active_context/databases/concerns/client.rb

+# frozen_string_literal: true
+
+module ActiveContext
+  module Databases
+    module Concerns
+      module Client
+        DEFAULT_PREFIX = 'gitlab'
+
+        attr_reader :options
+
+        def prefix
+          options[:prefix] || DEFAULT_PREFIX
+        end
+
+        def search(_)
+          raise NotImplementedError
+        end
+      end
+    end
+  end
+end

gems/gitlab-active-context/lib/active_context/databases/concerns/query_result.rb

+# frozen_string_literal: true
+
+module ActiveContext
+  module Databases
+    module Concerns
+      module QueryResult
+        include Enumerable
+
+        def each
+          raise NotImplementedError
+        end
+      end
+    end
+  end
+end

gems/gitlab-active-context/lib/active_context/databases/postgresql/adapter.rb

+# frozen_string_literal: true
+
+module ActiveContext
+  module Databases
+    module Postgresql
+      class Adapter
+        include ActiveContext::Databases::Concerns::Adapter
+
+        def client_klass
+          ActiveContext::Databases::Postgresql::Client
+        end
+      end
+    end
+  end
+end

gems/gitlab-active-context/lib/active_context/databases/postgresql/client.rb

+# frozen_string_literal: true
+
+module ActiveContext
+  module Databases
+    module Postgresql
+      class Client
+        include ActiveContext::Databases::Concerns::Client
+
+        DEFAULT_POOL_SIZE = 5
+        DEFAULT_POOL_TIMEOUT = 5
+
+        def initialize(options)
+          @options = options
+          @pool = ConnectionPool.new(
+            size: options.fetch(:pool_size, DEFAULT_POOL_SIZE),
+            timeout: options.fetch(:pool_timeout, DEFAULT_POOL_TIMEOUT)
+          ) do
+            PG.connect(connection_params)
+          end
+        end
+
+        def search(_query)
+          with_connection do |conn|
+            res = conn.exec('SELECT * FROM pg_stat_activity')
+            QueryResult.new(res)
+          end
+        end
+
+        private
+
+        def with_connection
+          @pool.with do |conn|
+            yield(conn)
+          end
+        end
+
+        def close
+          @pool&.shutdown(&:close)
+        end
+
+        def connection_params
+          {
+            host: options[:host],
+            port: options[:port],
+            dbname: options[:database],
+            user: options[:username],
+            password: options[:password],
+            connect_timeout: options.fetch(:connect_timeout, 5)
+          }.compact
+        end
+      end
+    end
+  end
+end

gems/gitlab-active-context/lib/active_context/databases/postgresql/query_result.rb

+# frozen_string_literal: true
+
+module ActiveContext
+  module Databases
+    module Postgresql
+      class QueryResult
+        include ActiveContext::Databases::Concerns::QueryResult
+
+        def initialize(pg_result)
+          @pg_result = pg_result
+        end
+
+        def each
+          return enum_for(:each) unless block_given?
+
+          pg_result.each do |row|
+            yield row
+          end
+        end
+
+        def count
+          pg_result.ntuples
+        end
+
+        def clear
+          pg_result.clear if pg_result.respond_to?(:clear)
+        end
+
+        private
+
+        attr_reader :pg_result
+      end
+    end
+  end
+end

gems/gitlab-active-context/spec/active_context_spec.rb

@@ -28,4 +28,35 @@
       expect(ActiveContext::Config.logger).to be_a(::Logger)
     end
   end
+
+  describe '.config' do
+    it 'returns the current configuration' do
+      config = described_class.config
+      expect(config).to be_a(ActiveContext::Config::Cfg)
+    end
+  end
+
+  describe '.adapter' do
+    it 'returns nil when not configured' do
+      expect(described_class.adapter).to be_nil
+    end
+
+    it 'returns configured adapter' do
+      described_class.configure do |config|
+        config.enabled = true
+        config.databases = {
+          main: {
+            adapter: 'ActiveContext::Databases::Postgresql::Adapter',
+            options: {
+              host: 'localhost',
+              port: 5432,
+              database: 'test_db'
+            }
+          }
+        }
+      end
+
+      expect(described_class.adapter).to be_a(ActiveContext::Databases::Postgresql::Adapter)
+    end
+  end
 end

gems/gitlab-active-context/spec/lib/active_context/databases/postgresql/adapter_spec.rb

+# frozen_string_literal: true
+
+RSpec.describe ActiveContext::Databases::Postgresql::Adapter do
+  let(:options) do
+    {
+      host: 'localhost',
+      port: 5432,
+      database: 'test_db',
+      username: 'user',
+      password: 'pass'
+    }
+  end
+
+  subject(:adapter) { described_class.new(options) }
+
+  it 'delegates search to client' do
+    query = ActiveContext::Query.filter(foo: :bar)
+    expect(adapter.client).to receive(:search).with(query)
+
+    adapter.search(query)
+  end
+end

gems/gitlab-active-context/spec/lib/active_context/databases/postgresql/client_spec.rb

+# frozen_string_literal: true
+
+RSpec.describe ActiveContext::Databases::Postgresql::Client do
+  let(:options) do
+    {
+      host: 'localhost',
+      port: 5432,
+      database: 'test_db',
+      username: 'user',
+      password: 'pass',
+      pool_size: 2,
+      pool_timeout: 1
+    }
+  end
+
+  subject(:client) { described_class.new(options) }
+
+  describe '#initialize' do
+    it 'creates a connection pool' do
+      expect(ConnectionPool).to receive(:new)
+        .with(hash_including(size: 2, timeout: 1))
+
+      client
+    end
+  end
+
+  describe '#search' do
+    let(:connection) { instance_double(PG::Connection) }
+    let(:query_result) { instance_double(PG::Result) }
+
+    before do
+      allow(PG).to receive(:connect).and_return(connection)
+      allow(connection).to receive(:exec).and_return(query_result)
+    end
+
+    it 'executes query and returns QueryResult' do
+      expect(connection).to receive(:exec).with('SELECT * FROM pg_stat_activity')
+      expect(ActiveContext::Databases::Postgresql::QueryResult)
+        .to receive(:new).with(query_result)
+
+      client.search('test query')
+    end
+  end
+
+  describe '#prefix' do
+    it 'returns default prefix when not specified' do
+      expect(client.prefix).to eq('gitlab')
+    end
+
+    it 'returns configured prefix' do
+      client = described_class.new(options.merge(prefix: 'custom'))
+      expect(client.prefix).to eq('custom')
+    end
+  end
+end

gems/gitlab-active-context/spec/lib/active_context/databases/postgresql/query_result_spec.rb

+# frozen_string_literal: true
+
+RSpec.describe ActiveContext::Databases::Postgresql::QueryResult do
+  let(:pg_result) { instance_double(PG::Result) }
+
+  subject(:query_result) { described_class.new(pg_result) }
+
+  describe '#each' do
+    it 'yields each row' do
+      rows = [
+        { 'id' => 1, 'name' => 'test1' },
+        { 'id' => 2, 'name' => 'test2' }
+      ]
+
+      allow(pg_result).to receive(:each).and_yield(rows[0]).and_yield(rows[1])
+
+      expect { |b| query_result.each(&b) }.to yield_successive_args(*rows)
+    end
+
+    it 'returns enumerator when no block given' do
+      expect(query_result.each).to be_a(Enumerator)
+    end
+  end
+
+  describe '#count' do
+    it 'returns number of tuples' do
+      allow(pg_result).to receive(:ntuples).and_return(5)
+      expect(query_result.count).to eq(5)
+    end
+  end
+
+  describe '#clear' do
+    context 'when pg_result responds to clear' do
+      before do
+        allow(pg_result).to receive(:respond_to?).with(:clear).and_return(true)
+      end
+
+      it 'clears the result' do
+        expect(pg_result).to receive(:clear)
+        query_result.clear
+      end
+    end
+
+    context 'when pg_result does not respond to clear' do
+      before do
+        allow(pg_result).to receive(:respond_to?).with(:clear).and_return(false)
+      end
+
+      it 'does nothing' do
+        expect(pg_result).not_to receive(:clear)
+        query_result.clear
+      end
+    end
+  end
+end