activerecord-import 0.14.1 → 1.4.1

data/lib/activerecord-import/adapters/postgresql_adapter.rb

@@ -1,12 +1,16 @@
+# frozen_string_literal: true
+
 module ActiveRecord::Import::PostgreSQLAdapter
   include ActiveRecord::Import::ImportSupport
   include ActiveRecord::Import::OnDuplicateKeyUpdateSupport
 
   MIN_VERSION_FOR_UPSERT = 90_500
 
-  def insert_many( sql, values, *args ) # :nodoc:
+  def insert_many( sql, values, options = {}, *args ) # :nodoc:
     number_of_inserts = 1
+    returned_values = {}
     ids = []
+    results = []
 
     base_sql, post_sql = if sql.is_a?( String )
       [sql, '']
@@ -15,15 +19,54 @@ module ActiveRecord::Import::PostgreSQLAdapter
     end
 
     sql2insert = base_sql + values.join( ',' ) + post_sql
-    if post_sql =~ /RETURNING\s/
-      ids = select_values( sql2insert, *args )
-    else
+
+    selections = returning_selections(options)
+    if selections.blank? || (options[:no_returning] && !options[:recursive])
       insert( sql2insert, *args )
+    else
+      returned_values = if selections.size > 1
+        # Select composite columns
+        db_result = select_all( sql2insert, *args )
+        { values: db_result.rows, columns: db_result.columns }
+      else
+        { values: select_values( sql2insert, *args ) }
+      end
+      clear_query_cache if query_cache_enabled
+    end
+
+    if options[:returning].blank?
+      ids = Array(returned_values[:values])
+    elsif options[:primary_key].blank?
+      options[:returning_columns] ||= returned_values[:columns]
+      results = Array(returned_values[:values])
+    else
+      # split primary key and returning columns
+      ids, results, options[:returning_columns] = split_ids_and_results(returned_values, options)
     end
 
-    ActiveRecord::Base.connection.query_cache.clear
+    ActiveRecord::Import::Result.new([], number_of_inserts, ids, results)
+  end
+
+  def split_ids_and_results( selections, options )
+    ids = []
+    returning_values = []
+
+    columns = Array(selections[:columns])
+    values = Array(selections[:values])
+    id_indexes = Array(options[:primary_key]).map { |key| columns.index(key) }
+    returning_columns = columns.reject.with_index { |_, index| id_indexes.include?(index) }
+    returning_indexes = returning_columns.map { |column| columns.index(column) }
+
+    values.each do |value|
+      value_array = Array(value)
+      ids << id_indexes.map { |index| value_array[index] }
+      returning_values << returning_indexes.map { |index| value_array[index] }
+    end
 
-    [number_of_inserts, ids]
+    ids.map!(&:first) if id_indexes.size == 1
+    returning_values.map!(&:first) if returning_columns.size == 1
+
+    [ids, returning_values, returning_columns]
   end
 
   def next_value_for_sequence(sequence_name)
@@ -31,10 +74,36 @@ module ActiveRecord::Import::PostgreSQLAdapter
   end
 
   def post_sql_statements( table_name, options ) # :nodoc:
-    if options[:no_returning] || options[:primary_key].blank?
-      super(table_name, options)
-    else
-      super(table_name, options) << "RETURNING #{options[:primary_key]}"
+    sql = []
+
+    if supports_on_duplicate_key_update?
+      # Options :recursive and :on_duplicate_key_ignore are mutually exclusive
+      if (options[:ignore] || options[:on_duplicate_key_ignore]) && !options[:on_duplicate_key_update] && !options[:recursive]
+        sql << sql_for_on_duplicate_key_ignore( table_name, options[:on_duplicate_key_ignore] )
+      end
+    elsif logger && options[:on_duplicate_key_ignore] && !options[:on_duplicate_key_update]
+      logger.warn "Ignoring on_duplicate_key_ignore because it is not supported by the database."
+    end
+
+    sql += super(table_name, options)
+
+    selections = returning_selections(options)
+    unless selections.blank? || (options[:no_returning] && !options[:recursive])
+      sql << " RETURNING #{selections.join(', ')}"
+    end
+
+    sql
+  end
+
+  def returning_selections(options)
+    selections = []
+    column_names = Array(options[:model].column_names)
+
+    selections += Array(options[:primary_key]) if options[:primary_key].present?
+    selections += Array(options[:returning]) if options[:returning].present?
+
+    selections.map do |selection|
+      column_names.include?(selection.to_s) ? "\"#{selection}\"" : selection
     end
   end
 
@@ -63,65 +132,75 @@ module ActiveRecord::Import::PostgreSQLAdapter
   # Returns a generated ON CONFLICT DO UPDATE statement given the passed
   # in +args+.
   def sql_for_on_duplicate_key_update( table_name, *args ) # :nodoc:
-    arg = args.first
+    arg, primary_key, locking_column = args
     arg = { columns: arg } if arg.is_a?( Array ) || arg.is_a?( String )
     return unless arg.is_a?( Hash )
 
-    sql = " ON CONFLICT "
+    sql = ' ON CONFLICT '.dup
     conflict_target = sql_for_conflict_target( arg )
 
     columns = arg.fetch( :columns, [] )
+    condition = arg[:condition]
     if columns.respond_to?( :empty? ) && columns.empty?
       return sql << "#{conflict_target}DO NOTHING"
     end
 
-    conflict_target ||= sql_for_default_conflict_target( table_name )
+    conflict_target ||= sql_for_default_conflict_target( table_name, primary_key )
     unless conflict_target
       raise ArgumentError, 'Expected :conflict_target or :constraint_name to be specified'
     end
 
     sql << "#{conflict_target}DO UPDATE SET "
     if columns.is_a?( Array )
-      sql << sql_for_on_duplicate_key_update_as_array( table_name, columns )
+      sql << sql_for_on_duplicate_key_update_as_array( table_name, locking_column, columns )
     elsif columns.is_a?( Hash )
-      sql << sql_for_on_duplicate_key_update_as_hash( table_name, columns )
+      sql << sql_for_on_duplicate_key_update_as_hash( table_name, locking_column, columns )
     elsif columns.is_a?( String )
       sql << columns
     else
       raise ArgumentError, 'Expected :columns to be an Array or Hash'
     end
+
+    sql << " WHERE #{condition}" if condition.present?
+
     sql
   end
 
-  def sql_for_on_duplicate_key_update_as_array( table_name, arr ) # :nodoc:
+  def sql_for_on_duplicate_key_update_as_array( table_name, locking_column, arr ) # :nodoc:
     results = arr.map do |column|
       qc = quote_column_name( column )
       "#{qc}=EXCLUDED.#{qc}"
     end
+    increment_locking_column!(table_name, results, locking_column)
     results.join( ',' )
   end
 
-  def sql_for_on_duplicate_key_update_as_hash( table_name, hsh ) # :nodoc:
+  def sql_for_on_duplicate_key_update_as_hash( table_name, locking_column, hsh ) # :nodoc:
     results = hsh.map do |column1, column2|
       qc1 = quote_column_name( column1 )
       qc2 = quote_column_name( column2 )
       "#{qc1}=EXCLUDED.#{qc2}"
     end
+    increment_locking_column!(table_name, results, locking_column)
     results.join( ',' )
   end
 
   def sql_for_conflict_target( args = {} )
     constraint_name = args[:constraint_name]
     conflict_target = args[:conflict_target]
-    if constraint_name
+    index_predicate = args[:index_predicate]
+    if constraint_name.present?
       "ON CONSTRAINT #{constraint_name} "
-    elsif conflict_target
-      '(' << Array( conflict_target ).join( ', ' ) << ') '
+    elsif conflict_target.present?
+      sql = '(' + Array( conflict_target ).reject( &:blank? ).join( ', ' ) + ') '
+      sql += "WHERE #{index_predicate} " if index_predicate
+      sql
     end
   end
 
-  def sql_for_default_conflict_target( table_name )
-    "(#{primary_key( table_name )}) "
+  def sql_for_default_conflict_target( table_name, primary_key )
+    conflict_target = Array(primary_key).join(', ')
+    "(#{conflict_target}) " if conflict_target.present?
   end
 
   # Return true if the statement is a duplicate key record error
@@ -129,15 +208,15 @@ module ActiveRecord::Import::PostgreSQLAdapter
     exception.is_a?(ActiveRecord::StatementInvalid) && exception.to_s.include?('duplicate key')
   end
 
-  def supports_on_duplicate_key_update?(current_version = postgresql_version)
-    current_version >= MIN_VERSION_FOR_UPSERT
+  def supports_on_duplicate_key_update?
+    database_version >= MIN_VERSION_FOR_UPSERT
   end
 
-  def supports_on_duplicate_key_ignore?(current_version = postgresql_version)
-    supports_on_duplicate_key_update?(current_version)
+  def supports_setting_primary_key_of_imported_objects?
+    true
   end
 
-  def support_setting_primary_key_of_imported_objects?
-    true
+  def database_version
+    defined?(postgresql_version) ? postgresql_version : super
   end
 end

data/lib/activerecord-import/adapters/sqlite3_adapter.rb

@@ -1,25 +1,28 @@
+# frozen_string_literal: true
+
 module ActiveRecord::Import::SQLite3Adapter
   include ActiveRecord::Import::ImportSupport
+  include ActiveRecord::Import::OnDuplicateKeyUpdateSupport
 
-  MIN_VERSION_FOR_IMPORT = "3.7.11".freeze
+  MIN_VERSION_FOR_IMPORT = "3.7.11"
+  MIN_VERSION_FOR_UPSERT = "3.24.0"
   SQLITE_LIMIT_COMPOUND_SELECT = 500
 
   # Override our conformance to ActiveRecord::Import::ImportSupport interface
   # to ensure that we only support import in supported version of SQLite.
   # Which INSERT statements with multiple value sets was introduced in 3.7.11.
-  def supports_import?(current_version = sqlite_version)
-    if current_version >= MIN_VERSION_FOR_IMPORT
-      true
-    else
-      false
-    end
+  def supports_import?
+    database_version >= MIN_VERSION_FOR_IMPORT
+  end
+
+  def supports_on_duplicate_key_update?
+    database_version >= MIN_VERSION_FOR_UPSERT
   end
 
   # +sql+ can be a single string or an array. If it is an array all
   # elements that are in position >= 1 will be appended to the final SQL.
-  def insert_many(sql, values, *args) # :nodoc:
+  def insert_many( sql, values, _options = {}, *args ) # :nodoc:
     number_of_inserts = 0
-    ids = []
 
     base_sql, post_sql = if sql.is_a?( String )
       [sql, '']
@@ -30,22 +33,142 @@ module ActiveRecord::Import::SQLite3Adapter
     value_sets = ::ActiveRecord::Import::ValueSetsRecordsParser.parse(values,
       max_records: SQLITE_LIMIT_COMPOUND_SELECT)
 
-    value_sets.each do |value_set|
-      number_of_inserts += 1
-      sql2insert = base_sql + value_set.join( ',' ) + post_sql
-      first_insert_id = insert( sql2insert, *args )
-      last_insert_id = first_insert_id + value_set.size - 1
-      ids.concat((first_insert_id..last_insert_id).to_a)
+    transaction(requires_new: true) do
+      value_sets.each do |value_set|
+        number_of_inserts += 1
+        sql2insert = base_sql + value_set.join( ',' ) + post_sql
+        insert( sql2insert, *args )
+      end
+    end
+
+    ActiveRecord::Import::Result.new([], number_of_inserts, [], [])
+  end
+
+  def pre_sql_statements( options )
+    sql = []
+    # Options :recursive and :on_duplicate_key_ignore are mutually exclusive
+    if !supports_on_duplicate_key_update? && (options[:ignore] || options[:on_duplicate_key_ignore])
+      sql << "OR IGNORE"
+    end
+    sql + super
+  end
+
+  def post_sql_statements( table_name, options ) # :nodoc:
+    sql = []
+
+    if supports_on_duplicate_key_update?
+      # Options :recursive and :on_duplicate_key_ignore are mutually exclusive
+      if (options[:ignore] || options[:on_duplicate_key_ignore]) && !options[:on_duplicate_key_update]
+        sql << sql_for_on_duplicate_key_ignore( options[:on_duplicate_key_ignore] )
+      end
     end
 
-    [number_of_inserts, ids]
+    sql + super
   end
 
   def next_value_for_sequence(sequence_name)
     %{nextval('#{sequence_name}')}
   end
 
-  def support_setting_primary_key_of_imported_objects?
-    true
+  # Add a column to be updated on duplicate key update
+  def add_column_for_on_duplicate_key_update( column, options = {} ) # :nodoc:
+    arg = options[:on_duplicate_key_update]
+    if arg.is_a?( Hash )
+      columns = arg.fetch( :columns ) { arg[:columns] = [] }
+      case columns
+      when Array then columns << column.to_sym unless columns.include?( column.to_sym )
+      when Hash then columns[column.to_sym] = column.to_sym
+      end
+    elsif arg.is_a?( Array )
+      arg << column.to_sym unless arg.include?( column.to_sym )
+    end
+  end
+
+  # Returns a generated ON CONFLICT DO NOTHING statement given the passed
+  # in +args+.
+  def sql_for_on_duplicate_key_ignore( *args ) # :nodoc:
+    arg = args.first
+    conflict_target = sql_for_conflict_target( arg ) if arg.is_a?( Hash )
+    " ON CONFLICT #{conflict_target}DO NOTHING"
+  end
+
+  # Returns a generated ON CONFLICT DO UPDATE statement given the passed
+  # in +args+.
+  def sql_for_on_duplicate_key_update( table_name, *args ) # :nodoc:
+    arg, primary_key, locking_column = args
+    arg = { columns: arg } if arg.is_a?( Array ) || arg.is_a?( String )
+    return unless arg.is_a?( Hash )
+
+    sql = ' ON CONFLICT '.dup
+    conflict_target = sql_for_conflict_target( arg )
+
+    columns = arg.fetch( :columns, [] )
+    condition = arg[:condition]
+    if columns.respond_to?( :empty? ) && columns.empty?
+      return sql << "#{conflict_target}DO NOTHING"
+    end
+
+    conflict_target ||= sql_for_default_conflict_target( primary_key )
+    unless conflict_target
+      raise ArgumentError, 'Expected :conflict_target to be specified'
+    end
+
+    sql << "#{conflict_target}DO UPDATE SET "
+    if columns.is_a?( Array )
+      sql << sql_for_on_duplicate_key_update_as_array( table_name, locking_column, columns )
+    elsif columns.is_a?( Hash )
+      sql << sql_for_on_duplicate_key_update_as_hash( table_name, locking_column, columns )
+    elsif columns.is_a?( String )
+      sql << columns
+    else
+      raise ArgumentError, 'Expected :columns to be an Array or Hash'
+    end
+
+    sql << " WHERE #{condition}" if condition.present?
+
+    sql
+  end
+
+  def sql_for_on_duplicate_key_update_as_array( table_name, locking_column, arr ) # :nodoc:
+    results = arr.map do |column|
+      qc = quote_column_name( column )
+      "#{qc}=EXCLUDED.#{qc}"
+    end
+    increment_locking_column!(table_name, results, locking_column)
+    results.join( ',' )
+  end
+
+  def sql_for_on_duplicate_key_update_as_hash( table_name, locking_column, hsh ) # :nodoc:
+    results = hsh.map do |column1, column2|
+      qc1 = quote_column_name( column1 )
+      qc2 = quote_column_name( column2 )
+      "#{qc1}=EXCLUDED.#{qc2}"
+    end
+    increment_locking_column!(table_name, results, locking_column)
+    results.join( ',' )
+  end
+
+  def sql_for_conflict_target( args = {} )
+    conflict_target = args[:conflict_target]
+    index_predicate = args[:index_predicate]
+    if conflict_target.present?
+      sql = '(' + Array( conflict_target ).reject( &:blank? ).join( ', ' ) + ') '
+      sql += "WHERE #{index_predicate} " if index_predicate
+      sql
+    end
+  end
+
+  def sql_for_default_conflict_target( primary_key )
+    conflict_target = Array(primary_key).join(', ')
+    "(#{conflict_target}) " if conflict_target.present?
+  end
+
+  # Return true if the statement is a duplicate key record error
+  def duplicate_key_update_error?(exception) # :nodoc:
+    exception.is_a?(ActiveRecord::StatementInvalid) && exception.to_s.include?('duplicate key')
+  end
+
+  def database_version
+    defined?(sqlite_version) ? sqlite_version : super
   end
 end

data/lib/activerecord-import/base.rb

@@ -1,9 +1,11 @@
+# frozen_string_literal: true
+
 require "pathname"
 require "active_record"
 require "active_record/version"
 
 module ActiveRecord::Import
-  ADAPTER_PATH = "activerecord-import/active_record/adapters".freeze
+  ADAPTER_PATH = "activerecord-import/active_record/adapters"
 
   def self.base_adapter(adapter)
     case adapter
@@ -11,24 +13,29 @@ module ActiveRecord::Import
     when 'mysql2spatial' then 'mysql2'
     when 'spatialite' then 'sqlite3'
     when 'postgresql_makara' then 'postgresql'
+    when 'makara_postgis' then 'postgresql'
     when 'postgis' then 'postgresql'
+    when 'cockroachdb' then 'postgresql'
     else adapter
     end
   end
 
   # Loads the import functionality for a specific database adapter
   def self.require_adapter(adapter)
-    require File.join(ADAPTER_PATH, "/abstract_adapter")
-    begin
-      require File.join(ADAPTER_PATH, "/#{base_adapter(adapter)}_adapter")
-    rescue LoadError
-      # fallback
-    end
+    require File.join(ADAPTER_PATH, "/#{base_adapter(adapter)}_adapter")
+  rescue LoadError
+    # fallback
   end
 
   # Loads the import functionality for the passed in ActiveRecord connection
   def self.load_from_connection_pool(connection_pool)
-    require_adapter connection_pool.spec.config[:adapter]
+    adapter =
+      if connection_pool.respond_to?(:db_config) # ActiveRecord >= 6.1
+        connection_pool.db_config.adapter
+      else
+        connection_pool.spec.config[:adapter]
+      end
+    require_adapter adapter
   end
 end