pg_party 1.0.1 → 1.4.0

data/lib/pg_party/cache.rb

@@ -6,32 +6,67 @@ module PgParty
   class Cache
     LOCK = Mutex.new
 
-    class << self
-      def clear!
-        LOCK.synchronize { store.clear }
+    def initialize
+      # automatically initialize a new hash when
+      # accessing an object id that doesn't exist
+      @store = Hash.new { |h, k| h[k] = { models: {}, partitions: nil, partitions_with_subpartitions: nil } }
+    end
+
+    def clear!
+      LOCK.synchronize { @store.clear }
+
+      nil
+    end
+
+    def fetch_model(key, child_table, &block)
+      return block.call unless caching_enabled?
+
+      LOCK.synchronize { fetch_value(@store[key][:models], child_table.to_sym, block) }
+    end
+
+    def fetch_partitions(key, include_subpartitions, &block)
+      return block.call unless caching_enabled?
+      sub_key = include_subpartitions ? :partitions_with_subpartitions : :partitions
+
+      LOCK.synchronize { fetch_value(@store[key], sub_key, block) }
+    end
+
+    private
+
+    def caching_enabled?
+      PgParty.config.caching
+    end
+
+    def fetch_value(subhash, key, block)
+      entry = subhash[key]
 
-        nil
+      if entry.nil? || entry.expired?
+        entry = Entry.new(block.call)
+        subhash[key] = entry
       end
 
-      def fetch_model(key, child_table, &block)
-        LOCK.synchronize do
-          store[key][:models][child_table.to_sym] ||= block.call
-        end
+      entry.value
+    end
+
+    class Entry
+      attr_reader :value
+
+      def initialize(value)
+        @value = value
+        @timestamp = Time.now
       end
 
-      def fetch_partitions(key, &block)
-        LOCK.synchronize do
-          store[key][:partitions] ||= block.call
-        end
+      def expired?
+        ttl.positive? && Time.now - @timestamp > ttl
       end
 
       private
 
-      def store
-        # automatically initialize a new hash when
-        # accessing an object id that doesn't exist
-        @store ||= Hash.new { |h, k| h[k] = { models: {}, partitions: nil } }
+      def ttl
+        PgParty.config.caching_ttl
       end
     end
+
+    private_constant :Entry
   end
 end

data/lib/pg_party/config.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+module PgParty
+  class Config
+    attr_accessor \
+      :caching,
+      :caching_ttl,
+      :schema_exclude_partitions,
+      :create_template_tables,
+      :create_with_primary_key,
+      :include_subpartitions_in_partition_list
+
+    def initialize
+      @caching = true
+      @caching_ttl = -1
+      @schema_exclude_partitions = true
+      @create_template_tables = true
+      @create_with_primary_key = false
+      @include_subpartitions_in_partition_list = false
+    end
+  end
+end

data/lib/pg_party/hacks/postgresql_database_tasks.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module PgParty
+  module Hacks
+    module PostgreSQLDatabaseTasks
+      def run_cmd(cmd, args, action)
+        if action != "dumping" || !PgParty.config.schema_exclude_partitions
+          return super
+        end
+
+        partitions = begin
+          ActiveRecord::Base.connection.select_values(
+            "SELECT DISTINCT inhrelid::regclass::text FROM pg_inherits"
+          )
+        rescue
+          []
+        end
+
+        excluded_tables = partitions.flat_map { |table| ["-T", "*.#{table}"] }
+
+        super(cmd, args + excluded_tables, action)
+      end
+    end
+  end
+end

data/lib/pg_party/model/hash_methods.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+require "pg_party/model_decorator"
+require "ruby2_keywords"
+
+module PgParty
+  module Model
+    module HashMethods
+      ruby2_keywords def create_partition(*args)
+        PgParty::ModelDecorator.new(self).create_hash_partition(*args)
+      end
+
+      ruby2_keywords def partition_key_in(*args)
+        PgParty::ModelDecorator.new(self).hash_partition_key_in(*args)
+      end
+    end
+  end
+end

data/lib/pg_party/model/list_methods.rb

@@ -1,15 +1,20 @@
 # frozen_string_literal: true
 
 require "pg_party/model_decorator"
+require "ruby2_keywords"
 
 module PgParty
   module Model
     module ListMethods
-      def create_partition(*args)
+      ruby2_keywords def create_partition(*args)
         PgParty::ModelDecorator.new(self).create_list_partition(*args)
       end
 
-      def partition_key_in(*args)
+      ruby2_keywords def create_default_partition(*args)
+        PgParty::ModelDecorator.new(self).create_default_partition(*args)
+      end
+
+      ruby2_keywords def partition_key_in(*args)
         PgParty::ModelDecorator.new(self).list_partition_key_in(*args)
       end
     end

data/lib/pg_party/model/methods.rb

@@ -5,12 +5,16 @@ require "pg_party/model_injector"
 module PgParty
   module Model
     module Methods
-      def range_partition_by(key=nil, &blk)
-        PgParty::ModelInjector.new(self, key || blk).inject_range_methods
+      def range_partition_by(*key, &blk)
+        PgParty::ModelInjector.new(self, *key, &blk).inject_range_methods
       end
 
-      def list_partition_by(key=nil, &blk)
-        PgParty::ModelInjector.new(self, key || blk).inject_list_methods
+      def list_partition_by(*key, &blk)
+        PgParty::ModelInjector.new(self, *key, &blk).inject_list_methods
+      end
+
+      def hash_partition_by(*key, &blk)
+        PgParty::ModelInjector.new(self, *key, &blk).inject_hash_methods
       end
 
       def partitioned?

data/lib/pg_party/model/range_methods.rb

@@ -1,15 +1,20 @@
 # frozen_string_literal: true
 
 require "pg_party/model_decorator"
+require "ruby2_keywords"
 
 module PgParty
   module Model
     module RangeMethods
-      def create_partition(*args)
+      ruby2_keywords def create_partition(*args)
         PgParty::ModelDecorator.new(self).create_range_partition(*args)
       end
 
-      def partition_key_in(*args)
+      ruby2_keywords def create_default_partition(*args)
+        PgParty::ModelDecorator.new(self).create_default_partition(*args)
+      end
+
+      ruby2_keywords def partition_key_in(*args)
         PgParty::ModelDecorator.new(self).range_partition_key_in(*args)
       end
     end

data/lib/pg_party/model/shared_methods.rb

@@ -6,15 +6,25 @@ module PgParty
   module Model
     module SharedMethods
       def reset_primary_key
-        PgParty::ModelDecorator.new(self).partition_primary_key
+        return base_class.primary_key if self != base_class
+
+        partitions = partitions(include_subpartitions: true)
+        return get_primary_key(base_class.name) if partitions.empty?
+
+        first_partition = partitions.detect { |p| !connection.table_partitioned?(p) }
+        raise 'No leaf partitions exist for this model. Create a partition to contain your data' unless first_partition
+
+        in_partition(first_partition).get_primary_key(base_class.name)
       end
 
       def table_exists?
-        PgParty::ModelDecorator.new(self).partition_table_exists?
+        target_table = partitions.first || table_name
+
+        connection.schema_cache.data_source_exists?(target_table)
       end
 
-      def partitions
-        PgParty::ModelDecorator.new(self).partitions
+      def partitions(*args)
+        PgParty::ModelDecorator.new(self).partitions(*args)
       end
 
       def in_partition(*args)

data/lib/pg_party/model_decorator.rb

@@ -1,27 +1,9 @@
 # frozen_string_literal: true
 
-require "pg_party/cache"
-
 module PgParty
   class ModelDecorator < SimpleDelegator
-    def partition_primary_key
-      if self != base_class
-        base_class.primary_key
-      elsif partition_name = partitions.first
-        in_partition(partition_name).get_primary_key(base_class.name)
-      else
-        get_primary_key(base_class.name)
-      end
-    end
-
-    def partition_table_exists?
-      target_table = partitions.first || table_name
-
-      connection.schema_cache.data_source_exists?(target_table)
-    end
-
     def in_partition(child_table_name)
-      PgParty::Cache.fetch_model(cache_key, child_table_name) do
+      PgParty.cache.fetch_model(cache_key, child_table_name) do
         Class.new(__getobj__) do
           self.table_name = child_table_name
 
@@ -41,6 +23,11 @@ module PgParty
           def self.new(*args, &blk)
             superclass.new(*args, &blk)
           end
+
+          # to avoid unnecessary db lookups
+          def self.partitions
+            []
+          end
         end
       end
     end
@@ -49,7 +36,7 @@ module PgParty
       if complex_partition_key
         complex_partition_key_query("(#{partition_key}) = (?)", value)
       else
-        where(current_arel_table[partition_key].eq(value))
+        where_partition_key(:eq, value)
       end
     end
 
@@ -61,9 +48,9 @@ module PgParty
           end_range
         )
       else
-        node = current_arel_table[partition_key]
-
-        where(node.gteq(start_range).and(node.lt(end_range)))
+        where_partition_key(:gteq, start_range).merge(
+          where_partition_key(:lt, end_range)
+        )
       end
     end
 
@@ -75,15 +62,11 @@ module PgParty
       end
     end
 
-    def partitions
-      PgParty::Cache.fetch_partitions(cache_key) do
-        connection.select_values(<<-SQL)
-          SELECT pg_inherits.inhrelid::regclass::text
-          FROM pg_tables
-          INNER JOIN pg_inherits
-            ON pg_tables.tablename::regclass = pg_inherits.inhparent::regclass
-          WHERE pg_tables.tablename = #{connection.quote(table_name)}
-        SQL
+    alias_method :hash_partition_key_in, :list_partition_key_in
+
+    def partitions(include_subpartitions: PgParty.config.include_subpartitions_in_partition_list)
+      PgParty.cache.fetch_partitions(cache_key, include_subpartitions) do
+        connection.partitions_for_table_name(table_name, include_subpartitions: include_subpartitions)
       end
     rescue
       []
@@ -108,6 +91,23 @@ module PgParty
       create_partition(:create_list_partition_of, table_name, **modified_options)
     end
 
+    def create_hash_partition(modulus:, remainder:, **options)
+      modified_options = options.merge(
+        modulus: modulus,
+        remainder: remainder,
+        primary_key: primary_key,
+      )
+
+      create_partition(:create_hash_partition_of, table_name, **modified_options)
+    end
+
+    def create_default_partition(**options)
+      modified_options = options.merge(
+        primary_key: primary_key,
+      )
+      create_partition(:create_default_partition_of, table_name, **modified_options)
+    end
+
     private
 
     def create_partition(migration_method, table_name, **options)
@@ -138,21 +138,33 @@ module PgParty
       end
     end
 
-    def model_class
-      if respond_to?(:klass)
-        klass
-      else
-        self
-      end
-    end
-
     def complex_partition_key_query(clause, *interpolated_values)
-      subquery = model_class
-        .unscoped
+      subquery = unscoped
         .select("*")
         .where(clause, *interpolated_values)
 
       from(subquery, current_alias)
     end
+
+    def where_partition_key(meth, values)
+      partition_key_array = Array.wrap(partition_key)
+      values = Array.wrap(values)
+
+      if partition_key_array.size != values.size
+        raise "number of provided values does not match the number of partition key columns"
+      end
+
+      arel_query = partition_key_array.zip(values).inject(nil) do |obj, (column, value)|
+        node = current_arel_table[column].send(meth, value)
+
+        if obj.nil?
+          node
+        else
+          obj.and(node)
+        end
+      end
+
+      where(arel_query)
+    end
   end
 end

data/lib/pg_party/model_injector.rb

@@ -2,9 +2,10 @@
 
 module PgParty
   class ModelInjector
-    def initialize(model, key)
+    def initialize(model, *key, &blk)
       @model = model
-      @key = key
+      @key = key.flatten.compact
+      @key_blk = blk
     end
 
     def inject_range_methods
@@ -19,6 +20,12 @@ module PgParty
       inject_methods_for(PgParty::Model::ListMethods)
     end
 
+    def inject_hash_methods
+      require "pg_party/model/hash_methods"
+
+      inject_methods_for(PgParty::Model::HashMethods)
+    end
+
     private
 
     def inject_methods_for(mod)
@@ -38,11 +45,16 @@ module PgParty
         instance_predicate: false
       )
 
-      if @key.is_a?(Proc)
-        @model.partition_key = @key.call
+      if @key_blk
+        @model.partition_key = @key_blk.call
         @model.complex_partition_key = true
       else
-        @model.partition_key = @key
+        if @key.size == 1
+          @model.partition_key = @key.first
+        else
+          @model.partition_key = @key
+        end
+
         @model.complex_partition_key = false
       end
     end

data/lib/pg_party/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module PgParty
-  VERSION = "1.0.1"
+  VERSION = "1.4.0"
 end