lib/wcc/contentful/store/postgres_store.rb in wcc-contentful-0.4.0.pre.alpha vs lib/wcc/contentful/store/postgres_store.rb in wcc-contentful-0.4.0.pre.beta

@@ -1,64 +1,69 @@
 # frozen_string_literal: true
 
 gem 'pg', '~> 1.0'
+gem 'connection_pool', '~> 2.2'
 require 'pg'
+require 'connection_pool'
 
 module WCC::Contentful::Store
   class PostgresStore < Base
-    def initialize(_config = nil, connection_options = nil)
+    attr_reader :connection_pool
+
+    def initialize(_config = nil, connection_options = nil, pool_options = nil)
       super()
+      @schema_ensured = false
       connection_options ||= { dbname: 'postgres' }
-      @conn = PG.connect(connection_options)
-      PostgresStore.ensure_schema(@conn)
+      pool_options ||= {}
+      @connection_pool = build_connection_pool(connection_options, pool_options)
     end
 
     def set(key, value)
       ensure_hash value
-      result = @conn.exec_prepared('upsert_entry', [key, value.to_json])
+      result = @connection_pool.with { |conn| conn.exec_prepared('upsert_entry', [key, value.to_json]) }
       return if result.num_tuples == 0
 
       val = result.getvalue(0, 0)
       JSON.parse(val) if val
     end
 
     def keys
-      result = @conn.exec_prepared('select_ids')
+      result = @connection_pool.with { |conn| conn.exec_prepared('select_ids') }
       arr = []
       result.each { |r| arr << r['id'].strip }
       arr
     end
 
     def delete(key)
-      result = @conn.exec_prepared('delete_by_id', [key])
+      result = @connection_pool.with { |conn| conn.exec_prepared('delete_by_id', [key]) }
       return if result.num_tuples == 0
 
       JSON.parse(result.getvalue(0, 1))
     end
 
     def find(key, **_options)
-      result = @conn.exec_prepared('select_entry', [key])
+      result = @connection_pool.with { |conn| conn.exec_prepared('select_entry', [key]) }
       return if result.num_tuples == 0
 
       JSON.parse(result.getvalue(0, 1))
     end
 
     def find_all(content_type:, options: nil)
       statement = "WHERE data->'sys'->'contentType'->'sys'->>'id' = $1"
       Query.new(
         self,
-        @conn,
+        @connection_pool,
         statement,
         [content_type],
         options
       )
     end
 
     class Query < Base::Query
-      def initialize(store, conn, statement = nil, params = nil, options = nil)
+      def initialize(store, connection_pool, statement = nil, params = nil, options = nil)
         super(store)
-        @conn = conn
+        @connection_pool = connection_pool
         @statement = statement ||
           "WHERE data->'sys'->>'id' IS NOT NULL"
         @params = params || []
         @options = options || {}
       end
@@ -72,30 +77,30 @@
         statement = @statement + " AND data->'fields'->$#{push_param(field, params)}" \
           "->$#{push_param(locale, params)} ? $#{push_param(expected, params)}"
 
         Query.new(
           @store,
-          @conn,
+          @connection_pool,
           statement,
           params,
           @options
         )
       end
 
       def count
         return @count if @count
 
         statement = 'SELECT count(*) FROM contentful_raw ' + @statement
-        result = @conn.exec(statement, @params)
+        result = @connection_pool.with { |conn| conn.exec(statement, @params) }
         @count = result.getvalue(0, 0).to_i
       end
 
       def first
         return @first if @first
 
         statement = 'SELECT * FROM contentful_raw ' + @statement + ' LIMIT 1'
-        result = @conn.exec(statement, @params)
+        result = @connection_pool.with { |conn| conn.exec(statement, @params) }
         return if result.num_tuples == 0
 
         resolve_includes(
           JSON.parse(result.getvalue(0, 1)),
           @options[:include]
@@ -133,11 +138,11 @@
 
       def resolve
         return @resolved if @resolved
 
         statement = 'SELECT * FROM contentful_raw ' + @statement
-        @resolved = @conn.exec(statement, @params)
+        @resolved = @connection_pool.with { |conn| conn.exec(statement, @params) }
       end
 
       def push_param(param, params)
         params << param
         params.length
@@ -166,13 +171,29 @@
            RETURN prev;
         END;
         $$ LANGUAGE 'plpgsql';
       HEREDOC
       )
+    end
 
+    def self.prepare_statements(conn)
       conn.prepare('upsert_entry', 'SELECT * FROM upsert_entry($1,$2)')
       conn.prepare('select_entry', 'SELECT * FROM contentful_raw WHERE id = $1')
       conn.prepare('select_ids', 'SELECT id FROM contentful_raw')
       conn.prepare('delete_by_id', 'DELETE FROM contentful_raw WHERE id = $1 RETURNING *')
+    end
+
+    private
+
+    def build_connection_pool(connection_options, pool_options)
+      ConnectionPool.new(pool_options) do
+        PG.connect(connection_options).tap do |conn|
+          unless @schema_ensured
+            PostgresStore.ensure_schema(conn)
+            @schema_ensured = true
+          end
+          PostgresStore.prepare_statements(conn)
+        end
+      end
     end
   end
 end