staging_table 0.1.0

data/lib/staging_table/bulk_inserter.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+module StagingTable
+  class BulkInserter
+    attr_reader :model, :batch_size
+
+    def initialize(model, batch_size: 1000)
+      @model = model
+      @batch_size = batch_size
+    end
+
+    def insert(records)
+      return if records.empty?
+
+      unless records.all? { |r| r.is_a?(Hash) }
+        raise RecordError, "All records must be hashes. If passing ActiveRecord objects, use Session#insert which normalizes them automatically."
+      end
+
+      columns = records.first.keys.map(&:to_s)
+      quoted_columns = columns.map { |c| connection.quote_column_name(c) }.join(", ")
+      quoted_table = connection.quote_table_name(model.table_name)
+
+      records.each_slice(batch_size) do |batch|
+        values_list = batch.map do |record|
+          "(" + columns.map { |col| quote(record.key?(col.to_sym) ? record[col.to_sym] : record[col]) }.join(", ") + ")"
+        end.join(", ")
+
+        sql = "INSERT INTO #{quoted_table} (#{quoted_columns}) VALUES #{values_list}"
+        connection.execute(sql)
+      end
+    end
+
+    private
+
+    def connection
+      model.connection
+    end
+
+    def quote(value)
+      connection.quote(value)
+    end
+  end
+end

data/lib/staging_table/configuration.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+module StagingTable
+  class Configuration
+    attr_accessor :default_batch_size, :default_transfer_strategy
+
+    def initialize
+      @default_batch_size = 1000
+      @default_transfer_strategy = :insert
+    end
+  end
+end

data/lib/staging_table/errors.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+module StagingTable
+  class Error < StandardError; end
+
+  # Raised when configuration options are invalid
+  class ConfigurationError < Error; end
+
+  # Raised when the database adapter is not supported
+  class AdapterError < Error; end
+
+  # Raised when staging table operations fail
+  class TableError < Error; end
+
+  # Raised when transfer strategy fails or is misconfigured
+  class TransferError < Error; end
+
+  # Raised when record data is invalid for insertion
+  class RecordError < Error; end
+end

data/lib/staging_table/instrumentation.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+require "active_support/notifications"
+
+module StagingTable
+  # Provides ActiveSupport::Notifications instrumentation for StagingTable operations.
+  #
+  # Available events:
+  #   - staging_table.create_table  - When a staging table is created
+  #   - staging_table.drop_table    - When a staging table is dropped
+  #   - staging_table.insert        - When records are inserted into staging
+  #   - staging_table.transfer      - When data is transferred to target table
+  #   - staging_table.stage         - Wraps the entire staging block operation
+  #
+  # Example:
+  #   ActiveSupport::Notifications.subscribe('staging_table.transfer') do |event|
+  #     Rails.logger.info "Transfer completed in #{event.duration}ms"
+  #     StatsD.measure('staging_table.transfer.duration', event.duration)
+  #   end
+  #
+  module Instrumentation
+    NAMESPACE = "staging_table"
+
+    EVENTS = %i[
+      create_table
+      drop_table
+      insert
+      transfer
+      stage
+    ].freeze
+
+    class << self
+      # Instruments a block with the given event name.
+      #
+      # @param event_name [Symbol] The event name (without namespace)
+      # @param payload [Hash] Additional payload data
+      # @yield The block to instrument
+      # @return The result of the block
+      def instrument(event_name, payload = {}, &block)
+        full_event_name = "#{NAMESPACE}.#{event_name}"
+        ActiveSupport::Notifications.instrument(full_event_name, payload, &block)
+      end
+
+      # Subscribe to a StagingTable event.
+      #
+      # @param event_name [Symbol, String] Event name (with or without namespace)
+      # @yield [event] Block called for each event
+      # @yieldparam event [ActiveSupport::Notifications::Event]
+      # @return [ActiveSupport::Notifications::Fanout::Subscribers::Evented]
+      def subscribe(event_name, &block)
+        full_name = event_name.to_s.start_with?(NAMESPACE) ? event_name : "#{NAMESPACE}.#{event_name}"
+        ActiveSupport::Notifications.subscribe(full_name, &block)
+      end
+
+      # Unsubscribe from a StagingTable event.
+      #
+      # @param subscriber [Object] The subscriber to remove
+      def unsubscribe(subscriber)
+        ActiveSupport::Notifications.unsubscribe(subscriber)
+      end
+
+      # Subscribe to all StagingTable events.
+      #
+      # @yield [event] Block called for each event
+      # @return [ActiveSupport::Notifications::Fanout::Subscribers::Evented]
+      def subscribe_all(&block)
+        ActiveSupport::Notifications.subscribe(/^#{NAMESPACE}\./o, &block)
+      end
+    end
+  end
+end

data/lib/staging_table/model_factory.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module StagingTable
+  class ModelFactory
+    def self.build(source_model, table_name, excluded_columns: [])
+      Class.new(source_model) do
+        self.table_name = table_name
+        self.ignored_columns = excluded_columns
+
+        # Ensure we don't inherit STI behavior unless intended for the temp table
+        self.inheritance_column = nil unless source_model.inheritance_column == "type" && source_model.columns_hash["type"]
+
+        def self.model_name
+          ActiveModel::Name.new(self, nil, superclass.name)
+        end
+
+        # Prevent the dynamic class from being added to the global constant namespace
+        def self.name
+          "#{superclass.name}::Staging_#{table_name}"
+        end
+      end
+    end
+  end
+end

data/lib/staging_table/session.rb

@@ -0,0 +1,186 @@
+# frozen_string_literal: true
+
+require "securerandom"
+
+module StagingTable
+  class Session
+    attr_reader :source_model, :staging_model, :options
+
+    # Supported callback options:
+    #   - before_insert: ->(session) { ... }
+    #   - after_insert: ->(session, records) { ... }
+    #   - before_transfer: ->(session) { ... }
+    #   - after_transfer: ->(session, result) { ... }
+    CALLBACK_OPTIONS = %i[before_insert after_insert before_transfer after_transfer].freeze
+
+    def initialize(source_model, **options)
+      @source_model = source_model
+      config = StagingTable.configuration
+      @callbacks = options.slice(*CALLBACK_OPTIONS)
+      @options = {
+        batch_size: config.default_batch_size,
+        transfer_strategy: config.default_transfer_strategy
+      }.merge(options.except(*CALLBACK_OPTIONS))
+      @table_created = false
+    end
+
+    def create_table
+      return if @table_created
+
+      payload = {
+        source_model: source_model,
+        source_table: source_model.table_name,
+        staging_table: staging_table_name
+      }
+
+      Instrumentation.instrument(:create_table, payload) do
+        adapter.create_table(staging_table_name, source_model.table_name, options)
+        @staging_model = ModelFactory.build(source_model, staging_table_name, excluded_columns: options[:excluded_columns] || [])
+        @table_created = true
+      end
+    end
+
+    def drop_table
+      return unless @table_created
+
+      payload = {
+        source_model: source_model,
+        source_table: source_model.table_name,
+        staging_table: staging_table_name
+      }
+
+      Instrumentation.instrument(:drop_table, payload) do
+        adapter.drop_table(staging_table_name)
+        @table_created = false
+        @staging_model = nil
+      end
+    end
+
+    def insert(records)
+      ensure_table_created!
+
+      run_callback(:before_insert, self)
+
+      normalized_records = normalize_records(records)
+
+      payload = {
+        source_model: source_model,
+        source_table: source_model.table_name,
+        staging_table: staging_table_name,
+        record_count: normalized_records.size,
+        batch_size: options[:batch_size] || 1000
+      }
+
+      Instrumentation.instrument(:insert, payload) do
+        BulkInserter.new(staging_model, batch_size: options[:batch_size] || 1000).insert(normalized_records)
+      end
+
+      run_callback(:after_insert, self, normalized_records)
+    end
+
+    def insert_from_query(relation)
+      ensure_table_created!
+
+      run_callback(:before_insert, self)
+
+      # TODO: Implement direct INSERT INTO SELECT for query-based insertion
+      # For now, we'll iterate, but this should be optimized
+      all_records = []
+
+      payload = {
+        source_model: source_model,
+        source_table: source_model.table_name,
+        staging_table: staging_table_name,
+        batch_size: options[:batch_size] || 1000
+      }
+
+      Instrumentation.instrument(:insert, payload) do |instrumentation_payload|
+        relation.find_in_batches(batch_size: options[:batch_size] || 1000) do |batch|
+          records = batch.map(&:attributes)
+          all_records.concat(records)
+          BulkInserter.new(staging_model, batch_size: options[:batch_size] || 1000).insert(records)
+        end
+        instrumentation_payload[:record_count] = all_records.size
+      end
+
+      run_callback(:after_insert, self, all_records)
+    end
+
+    def transfer
+      ensure_table_created!
+
+      run_callback(:before_transfer, self)
+
+      strategy_name = options[:transfer_strategy].to_s.camelize
+      begin
+        strategy_class = TransferStrategies.const_get(strategy_name)
+      rescue NameError
+        raise ConfigurationError, "Invalid transfer strategy: #{options[:transfer_strategy]}. Available strategies: insert, upsert."
+      end
+
+      payload = {
+        source_model: source_model,
+        source_table: source_model.table_name,
+        staging_table: staging_table_name,
+        strategy: options[:transfer_strategy],
+        staged_count: staging_model.count
+      }
+
+      result = Instrumentation.instrument(:transfer, payload) do |instrumentation_payload|
+        transfer_result = strategy_class.new(source_model, staging_model, options).transfer
+        instrumentation_payload[:result] = transfer_result
+        transfer_result
+      end
+
+      run_callback(:after_transfer, self, result)
+
+      result
+    end
+
+    # Delegate unknown methods to the staging model (e.g. for querying)
+    def method_missing(method, *args, &block)
+      if staging_model.respond_to?(method)
+        staging_model.send(method, *args, &block)
+      else
+        super
+      end
+    end
+
+    def respond_to_missing?(method, include_private = false)
+      staging_model.respond_to?(method, include_private) || super
+    end
+
+    private
+
+    def run_callback(name, *args)
+      callback = @callbacks[name]
+      return unless callback
+
+      callback.call(*args)
+    end
+
+    def adapter
+      @adapter ||= Adapters::Base.for(source_model.connection)
+    end
+
+    def staging_table_name
+      @staging_table_name ||= "staging_#{source_model.table_name}_#{SecureRandom.hex(8)}"
+    end
+
+    def ensure_table_created!
+      raise TableError, "Staging table has not been created. You must call #create_table or use StagingTable.stage with a block before inserting or transferring data." unless @table_created
+    end
+
+    def normalize_records(records)
+      if records.is_a?(ActiveRecord::Relation)
+        records.map(&:attributes)
+      elsif records.respond_to?(:to_a)
+        records.to_a.map do |record|
+          record.is_a?(ActiveRecord::Base) ? record.attributes : record
+        end
+      else
+        records
+      end
+    end
+  end
+end

data/lib/staging_table/transfer_result.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+module StagingTable
+  # Holds statistics about a transfer operation
+  class TransferResult
+    attr_reader :inserted, :updated, :skipped, :total
+
+    def initialize(inserted: 0, updated: 0, skipped: 0)
+      @inserted = inserted
+      @updated = updated
+      @skipped = skipped
+      @total = inserted + updated + skipped
+    end
+
+    def to_h
+      {
+        inserted: inserted,
+        updated: updated,
+        skipped: skipped,
+        total: total
+      }
+    end
+
+    def success?
+      inserted > 0 || updated > 0
+    end
+
+    def empty?
+      total.zero?
+    end
+
+    def inspect
+      "#<StagingTable::TransferResult inserted=#{inserted} updated=#{updated} skipped=#{skipped} total=#{total}>"
+    end
+  end
+end

data/lib/staging_table/transfer_strategies/insert.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module StagingTable
+  module TransferStrategies
+    class Insert
+      def initialize(source_model, staging_model, options = {})
+        @source_model = source_model
+        @staging_model = staging_model
+        @options = options
+        @connection = source_model.connection
+      end
+
+      def transfer
+        staged_count = @staging_model.count
+        return TransferResult.new if staged_count.zero?
+
+        columns = @staging_model.column_names.map { |c| @connection.quote_column_name(c) }.join(", ")
+        source_table = @connection.quote_table_name(@source_model.table_name)
+        staging_table = @connection.quote_table_name(@staging_model.table_name)
+
+        sql = <<~SQL
+          INSERT INTO #{source_table} (#{columns})
+          SELECT #{columns} FROM #{staging_table}
+        SQL
+
+        @connection.execute(sql)
+
+        # For plain INSERT, all staged records are inserted (assuming no constraint violations)
+        TransferResult.new(inserted: staged_count)
+      end
+    end
+  end
+end

data/lib/staging_table/transfer_strategies/upsert.rb

@@ -0,0 +1,159 @@
+# frozen_string_literal: true
+
+module StagingTable
+  module TransferStrategies
+    class Upsert
+      def initialize(source_model, staging_model, options = {})
+        @source_model = source_model
+        @staging_model = staging_model
+        @options = options
+        @connection = source_model.connection
+      end
+
+      def transfer
+        @staged_count = @staging_model.count
+        return TransferResult.new if @staged_count.zero?
+
+        adapter_name = @connection.adapter_name.downcase
+        case adapter_name
+        when /postgresql/
+          postgresql_upsert
+        when /mysql/
+          mysql_upsert
+        when /sqlite/
+          sqlite_upsert
+        else
+          raise AdapterError, "Upsert strategy not supported for adapter: #{adapter_name}. Supported adapters are PostgreSQL, MySQL, and SQLite."
+        end
+      end
+
+      private
+
+      def postgresql_upsert
+        conflict_target = Array(@options[:conflict_target])
+        if conflict_target.empty?
+          raise ConfigurationError, "PostgreSQL upsert requires :conflict_target option specifying the unique constraint columns. Example: transfer_strategy: :upsert, conflict_target: [:email]"
+        end
+
+        columns = column_names.map { |c| quote_column(c) }.join(", ")
+        conflict_target_sql = conflict_target.map { |c| quote_column(c) }.join(", ")
+        source_table = quote_table(@source_model.table_name)
+        staging_table = quote_table(@staging_model.table_name)
+
+        if @options[:conflict_action] == :ignore
+          # Use RETURNING to count actual inserts (rows where xmax = 0 are new inserts)
+          sql = "INSERT INTO #{source_table} (#{columns}) SELECT #{columns} FROM #{staging_table}"
+          sql += " ON CONFLICT (#{conflict_target_sql}) DO NOTHING"
+
+          count_before = @source_model.count
+          @connection.execute(sql)
+          count_after = @source_model.count
+
+          inserted = count_after - count_before
+          skipped = @staged_count - inserted
+          TransferResult.new(inserted: inserted, skipped: skipped)
+        else
+          updates = column_names.reject { |c| conflict_target.map(&:to_s).include?(c.to_s) || c == "id" }
+            .map { |c| "#{quote_column(c)} = EXCLUDED.#{quote_column(c)}" }.join(", ")
+
+          # Count existing records that match conflict target before upsert
+          count_before = @source_model.count
+          @connection.execute(
+            "INSERT INTO #{source_table} (#{columns}) SELECT #{columns} FROM #{staging_table} " \
+            "ON CONFLICT (#{conflict_target_sql}) DO UPDATE SET #{updates}"
+          )
+          count_after = @source_model.count
+
+          inserted = count_after - count_before
+          updated = @staged_count - inserted
+          TransferResult.new(inserted: inserted, updated: updated)
+        end
+      end
+
+      def mysql_upsert
+        columns = column_names.map { |c| quote_column(c) }.join(", ")
+        source_table = quote_table(@source_model.table_name)
+        staging_table = quote_table(@staging_model.table_name)
+
+        count_before = @source_model.count
+
+        if @options[:conflict_action] == :ignore
+          sql = "INSERT IGNORE INTO #{source_table} (#{columns}) SELECT #{columns} FROM #{staging_table}"
+          @connection.execute(sql)
+
+          count_after = @source_model.count
+          inserted = count_after - count_before
+          skipped = @staged_count - inserted
+          TransferResult.new(inserted: inserted, skipped: skipped)
+        else
+          sql = "INSERT INTO #{source_table} (#{columns}) SELECT #{columns} FROM #{staging_table}"
+          updates = column_names.reject { |c| c == "id" }
+            .map { |c| "#{quote_column(c)} = VALUES(#{quote_column(c)})" }.join(", ")
+          sql += " ON DUPLICATE KEY UPDATE #{updates}"
+          @connection.execute(sql)
+
+          count_after = @source_model.count
+          inserted = count_after - count_before
+          updated = @staged_count - inserted
+          TransferResult.new(inserted: inserted, updated: updated)
+        end
+      end
+
+      def sqlite_upsert
+        conflict_target = Array(@options[:conflict_target])
+        if conflict_target.empty?
+          raise ConfigurationError, "SQLite upsert requires :conflict_target option specifying the unique constraint columns. Example: transfer_strategy: :upsert, conflict_target: [:email]"
+        end
+
+        columns = column_names.map { |c| quote_column(c) }.join(", ")
+        source_table = quote_table(@source_model.table_name)
+        staging_table = quote_table(@staging_model.table_name)
+
+        count_before = @source_model.count
+
+        if @options[:conflict_action] == :ignore
+          sql = "INSERT OR IGNORE INTO #{source_table} (#{columns}) SELECT #{columns} FROM #{staging_table}"
+          @connection.execute(sql)
+
+          count_after = @source_model.count
+          inserted = count_after - count_before
+          skipped = @staged_count - inserted
+          TransferResult.new(inserted: inserted, skipped: skipped)
+        else
+          conflict_target_sql = conflict_target.map { |c| quote_column(c) }.join(", ")
+          updates = column_names.reject { |c| conflict_target.map(&:to_s).include?(c.to_s) || c == "id" }
+            .map { |c| "#{quote_column(c)} = excluded.#{quote_column(c)}" }.join(", ")
+
+          # Build individual upsert statements for each record
+          @staging_model.all.each do |record|
+            values_sql = column_names.map { |c| quote(record[c]) }.join(", ")
+            upsert_sql = "INSERT INTO #{source_table} (#{columns}) VALUES (#{values_sql}) " \
+                         "ON CONFLICT (#{conflict_target_sql}) DO UPDATE SET #{updates}"
+            @connection.execute(upsert_sql)
+          end
+
+          count_after = @source_model.count
+          inserted = count_after - count_before
+          updated = @staged_count - inserted
+          TransferResult.new(inserted: inserted, updated: updated)
+        end
+      end
+
+      def quote(value)
+        @connection.quote(value)
+      end
+
+      def column_names
+        @staging_model.column_names
+      end
+
+      def quote_column(name)
+        @connection.quote_column_name(name)
+      end
+
+      def quote_table(name)
+        @connection.quote_table_name(name)
+      end
+    end
+  end
+end

data/lib/staging_table/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module StagingTable
+  VERSION = "0.1.0"
+end

data/lib/staging_table.rb

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+require "active_record"
+require "staging_table/version"
+require "staging_table/errors"
+require "staging_table/configuration"
+require "staging_table/instrumentation"
+require "staging_table/transfer_result"
+require "staging_table/session"
+require "staging_table/model_factory"
+require "staging_table/bulk_inserter"
+require "staging_table/adapters/base"
+require "staging_table/adapters/postgresql"
+require "staging_table/adapters/mysql"
+require "staging_table/adapters/sqlite"
+require "staging_table/transfer_strategies/insert"
+require "staging_table/transfer_strategies/upsert"
+
+module StagingTable
+  class << self
+    def configuration
+      @configuration ||= Configuration.new
+    end
+
+    def configure
+      yield(configuration)
+    end
+
+    # Stage data for bulk import into a model's table.
+    #
+    # @param source_model [Class] The ActiveRecord model to stage data for
+    # @param options [Hash] Configuration options
+    # @option options [Integer] :batch_size Number of records per batch (default: 1000)
+    # @option options [Symbol] :transfer_strategy :insert or :upsert (default: :insert)
+    # @option options [Array<Symbol>] :conflict_target Columns for upsert conflict detection
+    # @option options [Symbol] :conflict_action :update or :ignore for upsert conflicts
+    # @option options [Proc] :before_insert Called before inserting into staging
+    # @option options [Proc] :after_insert Called after inserting into staging
+    # @option options [Proc] :before_transfer Called before transferring to target
+    # @option options [Proc] :after_transfer Called after transferring to target
+    #
+    # @yield [session] Block for staging operations
+    # @yieldparam session [Session] The staging session
+    # @return [TransferResult, Session] TransferResult when block given, Session otherwise
+    def stage(source_model, **options, &block)
+      session = Session.new(source_model, **options)
+
+      if block
+        payload = {
+          source_model: source_model,
+          source_table: source_model.table_name,
+          options: options.except(*Session::CALLBACK_OPTIONS)
+        }
+
+        Instrumentation.instrument(:stage, payload) do |instrumentation_payload|
+          session.create_table
+          yield(session)
+          result = session.transfer
+          instrumentation_payload[:result] = result
+          result
+        ensure
+          session.drop_table
+        end
+      else
+        session.create_table
+        session
+      end
+    end
+  end
+end

data/rbs_collection.yaml

@@ -0,0 +1,18 @@
+# Run `rbs collection install` to install required type definitions
+# See: https://github.com/ruby/rbs/blob/master/docs/collection.md
+
+sources:
+  - type: git
+    name: ruby/gem_rbs_collection
+    remote: https://github.com/ruby/gem_rbs_collection.git
+    revision: main
+    repo_dir: gems
+
+path: .gem_rbs_collection
+
+gems:
+  - name: activerecord
+  - name: activesupport
+  # Ignore gems without RBS definitions in the collection
+  - name: prism
+    ignore: true

data/sig/manifest.yaml

@@ -0,0 +1,5 @@
+# RBS manifest file
+# Defines gem dependencies for type checking
+dependencies:
+  - name: activerecord
+  - name: activesupport