lib/wcc/contentful/store/postgres_store.rb in wcc-contentful-0.2.2 vs lib/wcc/contentful/store/postgres_store.rb in wcc-contentful-0.3.0.pre.rc

@@ -3,20 +3,22 @@
 gem 'pg', '~> 1.0'
 require 'pg'
 
 module WCC::Contentful::Store
   class PostgresStore < Base
-    def initialize(connection_options = nil)
+    def initialize(_config = nil, connection_options = nil)
       super()
       connection_options ||= { dbname: 'postgres' }
       @conn = PG.connect(connection_options)
       PostgresStore.ensure_schema(@conn)
     end
 
     def set(key, value)
+      ensure_hash value
       result = @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
@@ -27,34 +29,40 @@
     end
 
     def delete(key)
       result = @conn.exec_prepared('delete_by_id', [key])
       return if result.num_tuples == 0
+
       JSON.parse(result.getvalue(0, 1))
     end
 
-    def find(key)
+    def find(key, **_options)
       result = @conn.exec_prepared('select_entry', [key])
       return if result.num_tuples == 0
+
       JSON.parse(result.getvalue(0, 1))
     end
 
-    def find_all(content_type:)
+    def find_all(content_type:, options: nil)
       statement = "WHERE data->'sys'->'contentType'->'sys'->>'id' = $1"
       Query.new(
+        self,
         @conn,
         statement,
-        [content_type]
+        [content_type],
+        options
       )
     end
 
     class Query < Base::Query
-      def initialize(conn, statement = nil, params = nil)
+      def initialize(store, conn, statement = nil, params = nil, options = nil)
+        super(store)
         @conn = conn
         @statement = statement ||
           "WHERE data->'sys'->>'id' IS NOT NULL"
         @params = params || []
+        @options = options || {}
       end
 
       def eq(field, expected, context = nil)
         locale = context[:locale] if context.present?
         locale ||= 'en-US'
@@ -63,46 +71,71 @@
 
         statement = @statement + " AND data->'fields'->$#{push_param(field, params)}" \
           "->$#{push_param(locale, params)} ? $#{push_param(expected, params)}"
 
         Query.new(
+          @store,
           @conn,
           statement,
-          params
+          params,
+          @options
         )
       end
 
       def count
         return @count if @count
+
         statement = 'SELECT count(*) FROM contentful_raw ' + @statement
         result = @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)
-        JSON.parse(result.getvalue(0, 1))
+        return if result.num_tuples == 0
+
+        resolve_includes(
+          JSON.parse(result.getvalue(0, 1)),
+          @options[:include]
+        )
       end
 
       def map
         arr = []
-        resolve.each { |row| arr << yield(JSON.parse(row['data'])) }
+        resolve.each do |row|
+          arr << yield(
+          resolve_includes(
+            JSON.parse(row['data']),
+            @options[:include]
+          )
+          )
+        end
         arr
       end
 
       def result
         arr = []
-        resolve.each { |row| arr << JSON.parse(row['data']) }
+        resolve.each do |row|
+          arr <<
+            resolve_includes(
+              JSON.parse(row['data']),
+              @options[:include]
+            )
+        end
         arr
       end
 
+      # TODO: override resolve_includes to make it more efficient
+
       private
 
       def resolve
         return @resolved if @resolved
+
         statement = 'SELECT * FROM contentful_raw ' + @statement
         @resolved = @conn.exec(statement, @params)
       end
 
       def push_param(param, params)
@@ -118,11 +151,11 @@
           data jsonb
         );
         CREATE INDEX IF NOT EXISTS contentful_raw_value_type ON contentful_raw ((data->'sys'->>'type'));
         CREATE INDEX IF NOT EXISTS contentful_raw_value_content_type ON contentful_raw ((data->'sys'->'contentType'->'sys'->>'id'));
 
-        DROP FUNCTION IF EXISTS "upsert_entry";
+        DROP FUNCTION IF EXISTS "upsert_entry"(_id varchar, _data jsonb);
         CREATE FUNCTION "upsert_entry"(_id varchar, _data jsonb) RETURNS jsonb AS $$
         DECLARE
           prev jsonb;
         BEGIN
           SELECT data FROM contentful_raw WHERE id = _id INTO prev;
@@ -131,10 +164,10 @@
               UPDATE
               SET data = _data;
            RETURN prev;
         END;
         $$ LANGUAGE 'plpgsql';
-HEREDOC
+      HEREDOC
       )
 
       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')