lib/trifle/stats/driver/postgres.rb in trifle-stats-1.2.0 vs lib/trifle/stats/driver/postgres.rb in trifle-stats-1.3.0

@@ -14,58 +14,71 @@
           @client = client
           @table_name = table_name
           @separator = '::'
         end
 
-        def inc(keys:, **values)
-          keys.map do |key|
-            pkey = key.join(separator)
+        def self.setup!(client = PG::Connection.new, table_name: 'trifle_stats')
+          client.exec("CREATE TABLE #{table_name} (key VARCHAR(255) PRIMARY KEY, data JSONB NOT NULL DEFAULT '{}'::jsonb)") # rubocop:disable Layout/LineLength
+        end
 
-            _inc_all(key: pkey, data: self.class.pack(hash: values))
+        def inc(keys:, **values)
+          data = self.class.pack(hash: values)
+          client.transaction do |c|
+            keys.map do |key|
+              pkey = key.join(separator)
+              c.exec(inc_query(key: pkey, data: data))
+            end
           end
         end
 
-        def _inc_all(key:, data:)
-          query = "INSERT INTO trifle_stats(key, data) VALUES ('#{key}', '#{data.to_json}') ON CONFLICT (key) DO UPDATE SET data = " + # rubocop:disable Layout/LineLength
-                  data.inject('to_jsonb(trifle_stats.data)') { |o, (k, v)| "jsonb_set(#{o}, '{#{k}}', (COALESCE(trifle_stats.data->>'#{k}', '0')::int + #{v})::text::jsonb)" } # rubocop:disable Layout/LineLength
-
-          client.exec(query)
+        def inc_query(key:, data:)
+          <<-SQL
+            INSERT INTO #{table_name} (key, data) VALUES ('#{key}', '#{data.to_json}')
+            ON CONFLICT (key) DO UPDATE SET data =
+            #{data.inject("to_jsonb(#{table_name}.data)") { |o, (k, v)| "jsonb_set(#{o}, '{#{k}}', (COALESCE(trifle_stats.data->>'#{k}', '0')::int + #{v})::text::jsonb)" }};
+          SQL
         end
 
         def set(keys:, **values)
-          keys.map do |key|
-            pkey = key.join(separator)
-
-            _set_all(key: pkey, data: self.class.pack(hash: values))
+          data = self.class.pack(hash: values)
+          client.transaction do |c|
+            keys.map do |key|
+              pkey = key.join(separator)
+              c.exec(set_query(key: pkey, data: data))
+            end
           end
         end
 
-        def _set_all(key:, data:)
-          query = "INSERT INTO trifle_stats(key, data) VALUES ('#{key}', '#{data.to_json}') ON CONFLICT (key) DO UPDATE SET data = " + # rubocop:disable Layout/LineLength
-                  data.inject('to_jsonb(trifle_stats.data)') { |o, (k, v)| "jsonb_set(#{o}, '{#{k}}', (#{v})::text::jsonb)" } # rubocop:disable Layout/LineLength
-
-          client.exec(query)
+        def set_query(key:, data:)
+          <<-SQL
+            INSERT INTO #{table_name} (key, data) VALUES ('#{key}', '#{data.to_json}')
+            ON CONFLICT (key) DO UPDATE SET data =
+            #{data.inject("to_jsonb(#{table_name}.data)") { |o, (k, v)| "jsonb_set(#{o}, '{#{k}}', (#{v})::text::jsonb)" }}
+          SQL
         end
 
         def get(keys:)
           pkeys = keys.map { |key| key.join(separator) }
-          data = _get_all(keys: pkeys)
-          map = data.inject({}) { |o, d| o.merge(d['key'] => d['data']) }
+          data = get_all(keys: pkeys)
+          map = data.inject({}) { |o, d| o.merge(d[:key] => d[:data]) }
 
-          pkeys.map { |pkey| self.class.unpack(hash: map[pkey]) || {} }
+          pkeys.map { |pkey| self.class.unpack(hash: map.fetch(pkey, {})) }
         end
 
-        def _get_all(keys:)
-          results = client.exec_params(
-            "SELECT * FROM #{table_name} WHERE key IN ('#{keys.join("', '")}');"
-          ).to_a
+        def get_all(keys:)
+          results = client.exec_params(get_query(keys: keys)).to_a
 
           results.map do |r|
-            r['data'] = JSON.parse(r['data'])
-            r
+            { key: r['key'], data: JSON.parse(r['data']) }
           rescue JSON::ParserError
-            r
+            { key: r['key'], data: {} }
           end
+        end
+
+        def get_query(keys:)
+          <<-SQL
+            SELECT * FROM #{table_name} WHERE key IN ('#{keys.join("', '")}');
+          SQL
         end
       end
     end
   end
 end