Capture HABTM tables

lib/polo/adapters/mysql.rb

@@ -4,7 +4,8 @@ class MySQL
       def on_duplicate_key_update(inserts, records)
         insert_and_record = inserts.zip(records)
         insert_and_record.map do |insert, record|
-          values_syntax = record.attributes.keys.map do |key|
+          attrs = record.is_a?(Hash) ? record.fetch(:values) : record.attributes
+          values_syntax = attrs.keys.map do |key|
             "`#{key}` = VALUES(`#{key}`)"
           end
 

lib/polo/adapters/postgres.rb

@@ -21,8 +21,13 @@ def on_duplicate_key_update(inserts, records)
       def ignore_transform(inserts, records)
         insert_and_record = inserts.zip(records)
         insert_and_record.map do |insert, record|
-          table_name = record.class.arel_table.name
-          id = record[:id]
+          if record.is_a?(Hash)
+            id = record.fetch(:values)[:id]
+            table_name = record.fetch(:table_name)
+          else
+            id = record[:id]
+            table_name = record.class.arel_table.name
+          end
           insert = insert.gsub(/VALUES \((.+)\)$/m, 'SELECT \\1')
           insert << " WHERE NOT EXISTS (SELECT 1 FROM #{table_name} WHERE id=#{id});"
         end

lib/polo/collector.rb

@@ -17,7 +17,7 @@ def collect
       unprepared_statement do
         ActiveSupport::Notifications.subscribed(collector, 'sql.active_record') do
           base_finder = @base_class.includes(@dependency_tree).where(@base_class.primary_key => @id)
-          collect_sql(@base_class, base_finder.to_sql)
+          collect_sql(klass: @base_class, sql: base_finder.to_sql)
           base_finder.to_a
         end
       end
@@ -35,22 +35,22 @@ def collect
     #
     def collector
       lambda do |name, start, finish, id, payload|
-        return unless payload[:name] =~ /^(.*) Load$/
-        begin
-          class_name = $1.constantize
-          sql = payload[:sql]
-          collect_sql(class_name, sql)
-        rescue ActiveRecord::StatementInvalid, NameError
-          # invalid table name (common when prefetching schemas)
+        sql = payload[:sql]
+        if payload[:name] =~ /^HABTM_.* Load$/
+          collect_sql(connection: @base_class.connection, sql: sql)
+        elsif payload[:name] =~ /^(.*) Load$/
+          begin
+            class_name = $1.constantize
+            collect_sql(klass: class_name, sql: sql)
+          rescue ActiveRecord::StatementInvalid, NameError
+            # invalid table name (common when prefetching schemas)
+          end
         end
       end
     end
 
-    def collect_sql(klass, sql)
-      @selects << {
-        klass: klass,
-        sql: sql
-      }
+    def collect_sql(select)
+      @selects << select
     end
 
     def unprepared_statement

lib/polo/sql_translator.rb

@@ -36,7 +36,11 @@ def records
 
     def inserts
       records.map do |record|
-        raw_sql(record)
+        if record.is_a?(Hash)
+          raw_sql_from_hash(record)
+        else
+          raw_sql_from_record(record)
+        end
       end
     end
 
@@ -47,12 +51,22 @@ def inserts
     # It will make use of the InsertManager class from the Arel gem to generate
     # insert statements
     #
-    def raw_sql(record)
+    def raw_sql_from_record(record)
       record.class.arel_table.create_insert.tap do |insert_manager|
         insert_manager.insert(insert_values(record))
       end.to_sql
     end
 
+    # Internal: Generates an insert SQL statement from a hash of values
+    def raw_sql_from_hash(hash)
+      connection = ActiveRecord::Base.connection
+      attributes = hash.fetch(:values)
+      table_name = connection.quote_table_name(hash.fetch(:table_name))
+      columns = attributes.keys.map{|k| connection.quote_column_name(k)}.join(", ")
+      value_placeholders = attributes.values.map{|v| "?" }.join(", ")
+      ActiveRecord::Base.send(:sanitize_sql_array, ["INSERT INTO #{table_name} (#{columns}) VALUES (#{value_placeholders})", *attributes.values])
+    end
+
     # Internal: Returns an object's attribute definitions along with
     # their set values (for Rails 3.x).
     #
@@ -88,7 +102,7 @@ def insert_values(record)
     #
     module ActiveRecordFive
       # Based on the codepath used in Rails 5
-      def raw_sql(record)
+      def raw_sql_from_record(record)
         values = record.send(:arel_attributes_with_values_for_create, record.class.column_names)
         model = record.class
         substitutes, binds = model.unscoped.substitute_values(values)

lib/polo/translator.rb

@@ -20,15 +20,22 @@ def translate
     end
 
     def instances
-      active_record_instances = @selects.flat_map do |select|
+      active_record_selects, raw_selects = @selects.partition{|s| s[:klass]}
+
+      active_record_instances = active_record_selects.flat_map do |select|
         select[:klass].find_by_sql(select[:sql]).to_a
       end
 
       if fields = @configuration.blacklist
         obfuscate!(active_record_instances, fields)
       end
 
-      active_record_instances
+      raw_instance_values = raw_selects.flat_map do |select|
+        table_name = select[:sql][/^SELECT .* FROM (?:"|`)([^"`]+)(?:"|`)/, 1]
+        select[:connection].select_all(select[:sql]).map { |values| {table_name: table_name, values: values} }
+      end
+
+      active_record_instances + raw_instance_values
     end
 
     private

spec/adapters/mysql_spec.rb

@@ -37,5 +37,15 @@
       translated_sql = adapter.on_duplicate_key_update(inserts, records)
       expect(translated_sql).to eq(insert_netto)
     end
+    it 'works for hash-values instead of ActiveRecord instances' do
+      insert_netto = [
+        %q{INSERT INTO "chefs" ("id", "name", "email") VALUES (1, 'Netto', '[email protected]') ON DUPLICATE KEY UPDATE `id` = VALUES(`id`), `name` = VALUES(`name`), `email` = VALUES(`email`)}
+      ]
+
+      inserts = translator.inserts
+      records = [{table_name: "chefs", values: { id: 1, name: "Netto", email: "[email protected]"}}]
+      translated_sql = adapter.on_duplicate_key_update(inserts, records)
+      expect(translated_sql).to eq(insert_netto)
+    end
   end
 end

spec/adapters/postgres_spec.rb

@@ -30,5 +30,14 @@
       translated_sql = adapter.ignore_transform(inserts, records)
       expect(translated_sql).to eq(insert_netto)
     end
+
+    it 'works for hash-values instead of ActiveRecord instances' do
+      insert_netto = [%q{INSERT INTO "chefs" ("id", "name", "email") SELECT 1, 'Netto', '[email protected]' WHERE NOT EXISTS (SELECT 1 FROM chefs WHERE id=1);}]
+
+      inserts = translator.inserts
+      records = [{table_name: "chefs", values: { id: 1, name: "Netto", email: "[email protected]"}}]
+      translated_sql = adapter.ignore_transform(inserts, records)
+      expect(translated_sql).to eq(insert_netto)
+    end
   end
 end

spec/polo_spec.rb

@@ -48,6 +48,21 @@
     expect(inserts).to include(two_cheeses)
   end
 
+  it 'generates inserts for HABTM relationships' do
+    ar_version = ActiveRecord::VERSION::STRING
+    skip("Not supported on ActiveRecord #{ar_version}") if ar_version < "4.1.0"
+    habtm_inserts = [
+      %q{INSERT INTO "recipes_tags" ("recipe_id", "tag_id") VALUES (2, 2)},
+      %q{INSERT INTO "tags" ("id", "name") VALUES (2, 'burgers')}
+    ]
+
+    inserts = Polo.explore(AR::Chef, 1, :recipes => :tags)
+
+    habtm_inserts.each do |habtm_insert|
+      expect(inserts).to include(habtm_insert)
+    end
+  end
+
   it 'generates inserts for many to many relationships' do
     many_to_many_inserts = [
       %q{INSERT INTO "recipes_ingredients" ("id", "recipe_id", "ingredient_id") VALUES (1, 1, 1)},

spec/support/activerecord_models.rb

@@ -3,13 +3,18 @@ class Recipe < ActiveRecord::Base
     belongs_to :chef
     has_many :recipes_ingredients
     has_many :ingredients, through: :recipes_ingredients
+    has_and_belongs_to_many :tags
 
     serialize :metadata, JSON
   end
 
   class Ingredient < ActiveRecord::Base
   end
 
+  class Tag < ActiveRecord::Base
+    has_and_belongs_to_many :recipes
+  end
+
   class RecipesIngredient < ActiveRecord::Base
     belongs_to :recipe
     belongs_to :ingredient

spec/support/factories.rb

@@ -5,11 +5,13 @@ def self.create_netto
       AR::Recipe.create(title: 'Turkey Sandwich', chef: netto).tap do |r|
         r.ingredients.create(name: 'Turkey', quantity: 'a lot')
         r.ingredients.create(name: 'Cheese', quantity: '1 slice')
+        r.tags.create(name: 'sandwiches')
       end
 
       AR::Recipe.create(title: 'Cheese Burger', chef: netto).tap do |r|
         r.ingredients.create(name: 'Patty', quantity: '1')
         r.ingredients.create(name: 'Cheese', quantity: '2 slices')
+        r.tags.create(name: 'burgers')
       end
     end
 

spec/support/schema.rb

@@ -39,4 +39,13 @@
   create_table :employees, force: true do |t|
     t.column :name, :string
   end
+
+  create_table :tags, force: true do |t|
+    t.column :name, :string
+  end
+
+  create_table :recipes_tags, id: false, force: true do |t|
+    t.column :recipe_id, :integer
+    t.column :tag_id, :integer
+  end
 end