Add support for polymorphic associations

.rubocop.yml

@@ -3,14 +3,23 @@ require: rubocop-rspec
 AllCops:
   TargetRubyVersion: 2.6
 
+Layout/LineLength:
+  Max: 120
+
 Layout/MultilineMethodCallIndentation:
   EnforcedStyle: indented_relative_to_receiver
 
-Metrics/LineLength:
-  Max: 120
+Layout/SpaceAroundMethodCallOperator:
+  Enabled: true
+
+Lint/RaiseException:
+  Enabled: true
+
+Lint/StructNewOverride:
+  Enabled: true
 
 Naming/PredicateName:
-  NameWhitelist:
+  AllowedMethods:
     - has_many
 
 Style/Documentation:
@@ -19,9 +28,21 @@ Style/Documentation:
 Style/EmptyMethod:
   EnforcedStyle: expanded
 
+Style/ExponentialNotation:
+  Enabled: true
+
 Style/FormatStringToken:
   EnforcedStyle: template
 
+Style/HashEachMethods:
+  Enabled: true
+
+Style/HashTransformKeys:
+  Enabled: true
+
+Style/HashTransformValues:
+  Enabled: true
+
 Style/PercentLiteralDelimiters:
   PreferredDelimiters:
     default: ()

.ruby-version

@@ -1 +1 @@
-2.6.5
+2.6.6

Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    table_saw (2.6.0)
+    table_saw (2.7.0)
       activerecord (>= 5.2)
       pg
       thor
@@ -48,7 +48,7 @@ GEM
     erubi (1.9.0)
     i18n (1.8.2)
       concurrent-ruby (~> 1.0)
-    jaro_winkler (1.5.2)
+    jaro_winkler (1.5.4)
     json (2.2.0)
     loofah (2.4.0)
       crass (~> 1.0.2)
@@ -58,10 +58,10 @@ GEM
     minitest (5.14.0)
     nokogiri (1.10.9)
       mini_portile2 (~> 2.4.0)
-    parallel (1.17.0)
-    parser (2.6.3.0)
+    parallel (1.19.1)
+    parser (2.7.1.2)
       ast (~> 2.4.0)
-    pg (1.2.2)
+    pg (1.2.3)
     pry (0.12.2)
       coderay (~> 1.1.0)
       method_source (~> 0.9.0)
@@ -81,6 +81,7 @@ GEM
       thor (>= 0.20.3, < 2.0)
     rainbow (3.0.0)
     rake (13.0.1)
+    rexml (3.2.4)
     rspec (3.9.0)
       rspec-core (~> 3.9.0)
       rspec-expectations (~> 3.9.0)
@@ -94,13 +95,14 @@ GEM
       diff-lcs (>= 1.2.0, < 2.0)
       rspec-support (~> 3.9.0)
     rspec-support (3.9.2)
-    rubocop (0.71.0)
+    rubocop (0.82.0)
       jaro_winkler (~> 1.5.1)
       parallel (~> 1.10)
-      parser (>= 2.6)
+      parser (>= 2.7.0.1)
       rainbow (>= 2.2.2, < 4.0)
+      rexml
       ruby-progressbar (~> 1.7)
-      unicode-display_width (>= 1.4.0, < 1.7)
+      unicode-display_width (>= 1.4.0, < 2.0)
     rubocop-rspec (1.33.0)
       rubocop (>= 0.60.0)
     ruby-progressbar (1.10.1)
@@ -116,7 +118,7 @@ GEM
     thread_safe (0.3.6)
     tzinfo (1.2.6)
       thread_safe (~> 0.1)
-    unicode-display_width (1.6.0)
+    unicode-display_width (1.7.0)
     zeitwerk (2.3.0)
 
 PLATFORMS

lib/table_saw/associations.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+module TableSaw
+  class Associations
+    attr_reader :manifest
+
+    def initialize(manifest)
+      @manifest = manifest
+    end
+
+    def belongs_to
+      @belongs_to ||= foreign_keys.each_with_object(Hash.new { |h, k| h[k] = Set.new }) do |fk, memo|
+        memo[fk.from_table].add(fk)
+      end
+    end
+
+    def has_many
+      @has_many ||= foreign_keys.each_with_object(Hash.new { |h, k| h[k] = Set.new }) do |fk, memo|
+        memo[fk.to_table].add(fk)
+      end
+    end
+
+    private
+
+    def foreign_keys
+      @foreign_keys ||= manifest_foreign_keys + schema_foreign_keys
+    end
+
+    def manifest_foreign_keys
+      manifest.foreign_keys.map do |fk|
+        TableSaw::ForeignKey.new(from_table: fk['from_table'], from_column: fk['from_column'],
+                                 to_table: fk['to_table'], to_column: fk['to_column'])
+      end
+    end
+
+    def schema_foreign_keys
+      TableSaw.information_schema.foreign_key_relationships.foreign_keys
+    end
+  end
+end

lib/table_saw/dependency_graph/belongs_to_directives.rb

@@ -3,42 +3,46 @@
 module TableSaw
   module DependencyGraph
     class BelongsToDirectives
-      attr_reader :directive
+      QUERY = <<~SQL
+        select distinct %{column} from %{table_name} where %{clause} and %{column} is not null and %{polymorphic}
+      SQL
 
-      def initialize(directive)
+      attr_reader :manifest, :directive
+
+      def initialize(manifest, directive)
+        @manifest = manifest
         @directive = directive
       end
 
       def call
-        associations.map do |from_column, to_table|
-          TableSaw::DependencyGraph::AddDirective.new(to_table, ids: ids[from_column], partial: directive.partial?)
+        associations.map do |fk|
+          TableSaw::DependencyGraph::AddDirective.new(fk.to_table, ids: ids[fk.column.primary_key],
+                                                                   partial: directive.partial?)
         end
       end
 
       private
 
       def associations
-        TableSaw.information_schema.belongs_to.fetch(directive.table_name, {})
+        manifest.associations.belongs_to.fetch(directive.table_name, Set.new)
       end
 
       def ids
-        @ids ||= associations.each_key.each_with_object({}) do |column, memo|
-          memo[column] = query_result(column).map { |row| row[column] }
+        @ids ||= associations.each_with_object({}) do |fk, memo|
+          memo[fk.column.primary_key] = query_result(fk).map { |row| row[fk.column.primary_key] }
         end
       end
 
       # rubocop:disable Metrics/AbcSize
-      def query_result(column)
+      def query_result(foreign_key)
         return [] unless directive.selectable?
 
         TableSaw::Connection.exec(
-          format(
-            'select distinct %{column} from %{table_name} where %{clause} and %{column} is not null',
-            primary_key: directive.primary_key, column: column, table_name: directive.table_name,
-            clause: TableSaw::Queries::SerializeSqlInClause.new(directive.table_name,
-                                                                directive.primary_key,
-                                                                directive.ids).call
-          )
+          format(QUERY, column: foreign_key.column.primary_key, table_name: directive.table_name,
+                        clause: TableSaw::Queries::SerializeSqlInClause.new(directive.table_name,
+                                                                            directive.primary_key,
+                                                                            directive.ids).call,
+                        polymorphic: foreign_key.type_condition)
         )
       end
       # rubocop:enable Metrics/AbcSize

lib/table_saw/dependency_graph/dump_table.rb

@@ -31,7 +31,7 @@ def fetch_associations(directive)
       private
 
       def fetch_belongs_to(directive)
-        TableSaw::DependencyGraph::BelongsToDirectives.new(directive).call
+        TableSaw::DependencyGraph::BelongsToDirectives.new(manifest, directive).call
       end
 
       def fetch_has_many(directive)

lib/table_saw/dependency_graph/has_many_directives.rb

@@ -3,6 +3,10 @@
 module TableSaw
   module DependencyGraph
     class HasManyDirectives
+      QUERY = <<~SQL
+        select %{primary_key} from %{table} where %{clause} and %{polymorphic}
+      SQL
+
       attr_reader :manifest, :directive
 
       def initialize(manifest, directive)
@@ -11,10 +15,10 @@ def initialize(manifest, directive)
       end
 
       def call
-        valid_associations.map do |table, column|
+        valid_associations.map do |fk|
           TableSaw::DependencyGraph::AddDirective.new(
-            table,
-            ids: query_result(table, column).map { |r| r[TableSaw.schema_cache.primary_keys(table)] },
+            fk.from_table,
+            ids: query_result(fk).map { |r| r[TableSaw.schema_cache.primary_keys(fk.from_table)] },
             partial: directive.partial?
           )
         end
@@ -23,31 +27,32 @@ def call
       private
 
       def associations
-        TableSaw.information_schema.has_many.fetch(directive.table_name, [])
+        manifest.associations.has_many.fetch(directive.table_name, Set.new)
       end
 
       # rubocop:disable Metrics/AbcSize
       def valid_associations
-        associations.select do |table, _column|
-          next false if directive.partial? && TableSaw.schema_cache.primary_keys(table).nil?
-          next true if directive.has_many.include?(table)
+        associations.select do |fk|
+          next false if directive.partial? && TableSaw.schema_cache.primary_keys(fk.from_table).nil?
+          next true if directive.has_many.include?(fk.from_table)
 
-          manifest.has_many.fetch(directive.table_name, []).include?(table)
+          manifest.has_many.fetch(directive.table_name, []).include?(fk.from_table)
         end
       end
-      # rubocop:enable Metrics/AbcSize
 
-      def query_result(table, column)
+      def query_result(foreign_key)
         return [] unless directive.selectable?
 
         TableSaw::Connection.exec(
-          format(
-            'select %{primary_key} from %{table} where %{clause}',
-            primary_key: TableSaw.schema_cache.primary_keys(table), table: table,
-            clause: TableSaw::Queries::SerializeSqlInClause.new(table, column, directive.ids).call
-          )
+          format(QUERY, primary_key: TableSaw.schema_cache.primary_keys(foreign_key.from_table),
+                        table: foreign_key.from_table,
+                        clause: TableSaw::Queries::SerializeSqlInClause.new(foreign_key.from_table,
+                                                                            foreign_key.column.primary_key,
+                                                                            directive.ids).call,
+                        polymorphic: foreign_key.type_condition)
         )
       end
+      # rubocop:enable Metrics/AbcSize
     end
   end
 end

lib/table_saw/foreign_key.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+module TableSaw
+  class ForeignKey
+    class Column
+      REGEX = /(\w+)(?::(\w+)\((\w+)\))?/.freeze
+
+      attr_reader :value
+
+      def initialize(value)
+        @value = value
+      end
+
+      def primary_key
+        value[REGEX, 1]
+      end
+
+      def type_condition
+        polymorphic? ? "#{type_column} = '#{type_value}'" : '1 = 1'
+      end
+
+      private
+
+      def type_column
+        value[REGEX, 2]
+      end
+
+      def type_value
+        value[REGEX, 3]
+      end
+
+      def polymorphic?
+        !(type_column.nil? || type_value.nil?)
+      end
+    end
+
+    attr_reader :name, :from_table, :from_column, :to_table, :to_column
+
+    def initialize(name: nil, from_table:, from_column:, to_table:, to_column:)
+      @name = name
+      @from_table = from_table
+      @from_column = from_column
+      @to_table = to_table
+      @to_column = to_column
+    end
+
+    def type_condition
+      @type_condition ||= column.type_condition
+    end
+
+    def column
+      @column ||= Column.new(from_column)
+    end
+
+    def eql?(other)
+      hash == other.hash
+    end
+
+    def hash
+      [from_table, from_column, to_table, to_column].hash
+    end
+  end
+end

lib/table_saw/information_schema.rb

@@ -2,14 +2,10 @@
 
 module TableSaw
   class InformationSchema
-    %i(belongs_to constraint_names has_many).each do |method_name|
-      define_method method_name do
-        foreign_key_relationships.public_send method_name
-      end
+    def constraint_names
+      foreign_key_relationships.constraint_names
     end
 
-    private
-
     def foreign_key_relationships
       @foreign_key_relationships ||= TableSaw::Queries::ForeignKeyRelationships.new
     end