combinaut_stagehand 1.2.1

data/lib/stagehand/database.rb

@@ -0,0 +1,197 @@
+require 'thread'
+require 'stagehand/active_record_extensions'
+
+module Stagehand
+  module Database
+    extend self
+
+    def transaction
+      raise InvalidConnectionError, "Calling Stagehand::Database.transaction is not valid unless connected to staging" unless connected_to_staging?
+
+      success = false
+      attempts = 0
+      output = nil
+      ActiveRecord::Base.transaction do
+        Production::Record.transaction do
+          attempts += 1
+
+          raise NoRetryError, "Retrying is not allowed in Stagehand::Database.transaction" if attempts > 1
+
+          output = yield
+
+          success = true
+        end
+
+        raise ActiveRecord::Rollback unless success
+      end
+
+      return output
+    ensure
+      Rails.logger.warn "Stagehand::Database transaction was rolled back" unless success
+    end
+
+    def each(&block)
+      with_production_connection(&block) unless Configuration.single_connection?
+      with_staging_connection(&block)
+    end
+
+    def connected_to_production?
+      current_connection_name == Configuration.production_connection_name
+    end
+
+    def connected_to_staging?
+      current_connection_name == Configuration.staging_connection_name
+    end
+
+    def production_connection
+      ProductionProbe.connection
+    end
+
+    def staging_connection
+      StagingProbe.connection
+    end
+
+    def production_database_name
+      database_name(Configuration.production_connection_name)
+    end
+
+    def staging_database_name
+      database_name(Configuration.staging_connection_name)
+    end
+
+    def staging_database_versions
+      Stagehand::Database.staging_connection.select_values(versions_scope)
+    end
+
+    def production_database_versions
+      Stagehand::Database.production_connection.select_values(versions_scope)
+    end
+
+    def with_staging_connection(&block)
+      with_connection(Configuration.staging_connection_name, &block)
+    end
+
+    def with_production_connection(&block)
+      with_connection(Configuration.production_connection_name, &block)
+    end
+
+    def with_connection(connection_name, &block)
+      if current_connection_name != connection_name.to_sym
+        Rails.logger.debug "Connecting to #{connection_name}"
+        output = swap_connection(connection_name, &block)
+        Rails.logger.debug "Restoring connection to #{current_connection_name}"
+      else
+        Rails.logger.debug "Already connected to #{connection_name}"
+        output = yield connection_name
+      end
+      return output
+    end
+
+    private
+
+    def swap_connection(connection_name)
+      pushed = ConnectionStack.push(connection_name.to_sym)
+      cache = ActiveRecord::Base.connection_pool.query_cache_enabled
+      ActiveRecord::Base.connection_specification_name = current_connection_name
+      ActiveRecord::Base.connection_pool.enable_query_cache! if cache
+
+      yield connection_name
+    ensure
+      ConnectionStack.pop if pushed
+      ActiveRecord::Base.connection_specification_name = current_connection_name
+      ActiveRecord::Base.connection_pool.enable_query_cache! if cache
+    end
+
+    def current_connection_name
+      ConnectionStack.last
+    end
+
+    def database_name(connection_name)
+      database_configuration.dig(connection_name.to_s, 'database')
+    end
+
+    def database_configuration
+      @database_configuration ||= Rails.configuration.database_configuration
+    end
+
+    def versions_scope
+      ActiveRecord::SchemaMigration.order(:version)
+    end
+
+    # CLASSES
+
+    class Probe < ActiveRecord::Base
+      self.abstract_class = true
+      self.stagehand_threadsafe_connections = false # We don't want to track connection per-thread for Probes
+
+      # We fake the class name so we can create a connection pool with the desired connection name instead of the name of the class
+      def self.init_connection(connection_name)
+        @probe_name = connection_name
+        establish_connection(connection_name)
+      ensure
+        @probe_name = nil
+      end
+
+      def self.name
+        @probe_name || super
+      end
+    end
+
+    class StagingProbe < Probe
+      self.abstract_class = true
+
+      def self.init_connection
+        super(Configuration.staging_connection_name)
+      end
+
+      def self.connection
+        if Stagehand::Database.connected_to_staging?
+          ActiveRecord::Base.connection # Reuse existing connection so we stay within the current transaction
+        else
+          super
+        end
+      end
+
+      init_connection
+    end
+
+    class ProductionProbe < Probe
+      self.abstract_class = true
+
+      def self.init_connection
+        super(Configuration.production_connection_name)
+      end
+
+      init_connection unless Configuration.single_connection?
+    end
+
+    # Threadsafe tracking of the connection stack
+    module ConnectionStack
+      def self.push(connection_name)
+        current_stack.push connection_name
+      end
+
+      def self.pop
+        current_stack.pop
+      end
+
+      def self.last
+        current_stack.last
+      end
+
+      def self.current_stack
+        if stack = Thread.current.thread_variable_get('sparkle_connection_name_stack')
+          stack
+        else
+          stack = Concurrent::Array.new
+          stack << Rails.env.to_sym
+          Thread.current.thread_variable_set('sparkle_connection_name_stack', stack)
+          stack
+        end
+      end
+    end
+
+    class InvalidConnectionError < StandardError; end
+    class NoRetryError < StandardError; end
+  end
+end

data/lib/stagehand/engine.rb

@@ -0,0 +1,26 @@
+require 'stagehand/configuration'
+
+module Stagehand
+  class Engine < ::Rails::Engine
+    isolate_namespace Stagehand
+
+    config.generators do |g|
+      g.test_framework :rspec
+    end
+
+    # These require the rails application to be initialized because configuration variables are used
+    initializer 'stagehand.load_modules' do
+      require 'stagehand/cache'
+      require 'stagehand/key'
+      require 'stagehand/database'
+      require 'stagehand/connection_adapter_extensions'
+      require 'stagehand/controller_extensions'
+      require 'stagehand/active_record_extensions'
+      require 'stagehand/schema_extensions'
+      require 'stagehand/staging'
+      require 'stagehand/production'
+      require 'stagehand/schema'
+      require 'stagehand/auditor'
+    end
+  end
+end

data/lib/stagehand/key.rb

@@ -0,0 +1,23 @@
+module Stagehand
+  module Key
+    extend self
+
+    def generate(staging_record, options = {})
+      case staging_record
+      when Staging::CommitEntry
+        id = staging_record.record_id || staging_record.id
+        table_name = staging_record.table_name || staging_record.class.table_name
+      when ActiveRecord::Base
+        id = staging_record.id
+        table_name = staging_record.class.table_name
+      else
+        id = staging_record
+        table_name = options[:table_name]
+      end
+
+      raise 'Invalid input' unless table_name && id
+
+      return [table_name, id]
+    end
+  end
+end

data/lib/stagehand/production/controller.rb

@@ -0,0 +1,12 @@
+module Stagehand
+  module Production
+    module Controller
+      extend ActiveSupport::Concern
+      include Stagehand::ControllerExtensions
+
+      included do
+        use_production_database
+      end
+    end
+  end
+end

data/lib/stagehand/production.rb

@@ -0,0 +1,132 @@
+require 'stagehand/production/controller'
+
+module Stagehand
+  module Production
+    extend self
+
+    # Outputs a symbol representing the status of the staging record as it exists in the production database
+    def status(staging_record, table_name = nil)
+      if !exists?(staging_record, table_name)
+        :new
+      elsif modified?(staging_record, table_name)
+        :modified
+      else
+        :not_modified
+      end
+    end
+
+    def save(staging_record, table_name = nil)
+      attributes = staging_record_attributes(staging_record, table_name)
+
+      return unless attributes.present?
+
+      write(staging_record, attributes, table_name)
+    end
+
+    def write(staging_record, attributes, table_name = nil)
+      table_name, id = Stagehand::Key.generate(staging_record, :table_name => table_name)
+
+      production_record = Connection.with_production_writes do
+        prepare_to_modify(table_name)
+
+        if update(table_name, id, attributes).nonzero?
+          Record.find(id)
+        else
+          Record.find(insert(table_name, attributes))
+        end
+      end
+
+      return production_record
+    end
+
+    def delete(staging_record, table_name = nil)
+      Connection.with_production_writes do
+        matching(staging_record, table_name).delete_all
+      end
+    end
+
+    def exists?(staging_record, table_name = nil)
+      matching(staging_record, table_name).exists?
+    end
+
+    # Returns true if the staging record's attributes are different from the production record's attributes
+    # Returns true if the staging_record does not exist on production
+    # Returns false if the staging record is identical to the production record
+    def modified?(staging_record, table_name = nil)
+      production_record_attributes(staging_record, table_name) != staging_record_attributes(staging_record, table_name)
+    end
+
+    def find(*args)
+      matching(*args).first
+    end
+
+    # Returns a scope that limits results any occurrences of the specified record.
+    # Record can be specified by passing a staging record, or an id and table_name.
+    def matching(staging_record, table_name = nil)
+      table_name, id = Stagehand::Key.generate(staging_record, :table_name => table_name)
+      prepare_to_modify(table_name)
+      return Record.where(:id => id)
+    end
+
+    private
+
+    def production_record_attributes(staging_record, table_name = nil)
+      Record.connection.select_one(matching(staging_record, table_name))
+    end
+
+    def staging_record_attributes(staging_record, table_name = nil)
+      table_name, id = Stagehand::Key.generate(staging_record, :table_name => table_name)
+      hash = select(table_name, id)
+      hash.except(*ignored_columns(table_name)) if hash
+    end
+
+    def ignored_columns(table_name)
+      Array.wrap(Configuration.ignored_columns[table_name]).map(&:to_s)
+    end
+
+    def select(table_name, id)
+      table = Arel::Table.new(table_name)
+      statement = Arel::SelectManager.new
+      statement.from table
+      statement.project Arel.star
+      statement.where table[:id].eq(id)
+
+      Stagehand::Database::StagingProbe.connection.select_one(statement)
+    end
+
+    def update(table_name, id, attributes)
+      table = Arel::Table.new(table_name)
+      statement = Arel::UpdateManager.new
+      statement.table table
+      statement.set attributes.map {|attribute, value| [table[attribute], value] }
+      statement.where table[:id].eq(id)
+
+      Record.connection.update(statement)
+    end
+
+    def insert(table_name, attributes)
+      table = Arel::Table.new(table_name)
+      statement = Arel::InsertManager.new
+      statement.into table
+      statement.insert attributes.map {|attribute, value| [table[attribute], value] }
+
+      Record.connection.insert(statement)
+    end
+
+    def prepare_to_modify(table_name)
+      raise "Can't prepare to modify production records without knowning the table_name" unless table_name.present?
+
+      return if Record.table_name == table_name
+
+      Record.table_name = table_name
+      Record.reset_column_information
+    end
+
+    # CLASSES
+
+    class Record < Stagehand::Database::ProductionProbe
+      self.record_timestamps = false
+      self.inheritance_column = nil
+    end
+  end
+end

data/lib/stagehand/schema/statements.rb

@@ -0,0 +1,53 @@
+module Stagehand
+  module Schema
+    module Statements
+      # Ensure that developers are aware they need to make a determination of whether stagehand should track this table or not
+      def create_table(table_name, options = {})
+        super
+        options = options.symbolize_keys
+
+        return if Database.connected_to_production? && !Stagehand::Configuration.single_connection?
+        return if options[:stagehand] == false
+        return if UNTRACKED_TABLES.include?(table_name)
+
+        Schema.add_stagehand! :only => table_name
+      end
+
+      def rename_table(old_table_name, new_table_name, *)
+        return super unless Schema.has_stagehand?(old_table_name)
+
+        Schema.remove_stagehand!(:only => old_table_name, :remove_entries => false)
+        super
+        Schema.add_stagehand!(:only => new_table_name)
+        Staging::CommitEntry.where(:table_name => old_table_name).update_all(:table_name => new_table_name)
+      end
+
+      def drop_table(table_name, *)
+        return super unless Schema.has_stagehand?(table_name) && table_exists?(Staging::CommitEntry.table_name)
+
+        Schema.remove_stagehand!(:only => table_name, :remove_entries => true)
+        super
+      end
+    end
+
+    # Allow dumping of stagehand create_table directive
+    # e.g. create_table "comments", stagehand: true do |t|
+    module DumperExtensions
+      def table(table_name, stream)
+        stagehand = Stagehand::Schema.has_stagehand?(table_name)
+        stagehand = ':commit_entries' if table_name == Staging::CommitEntry.table_name
+
+        table_stream = StringIO.new
+        super(table_name, table_stream)
+        table_stream.rewind
+        table_schema = table_stream.read.gsub(/create_table (.+) do/, 'create_table \1' + ", stagehand: #{stagehand} do")
+        stream.puts table_schema
+
+        return stream
+      end
+    end
+  end
+end
+
+ActiveRecord::Base.connection.class.include Stagehand::Schema::Statements
+ActiveRecord::SchemaDumper.prepend Stagehand::Schema::DumperExtensions

data/lib/stagehand/schema.rb

@@ -0,0 +1,145 @@
+require "stagehand/schema/statements"
+
+module Stagehand
+  module Schema
+    extend self
+
+    UNTRACKED_TABLES = ['ar_internal_metadata', 'schema_migrations', 'data_migrations', Stagehand::Staging::CommitEntry.table_name]
+
+    def init_stagehand!(**table_options)
+      ActiveRecord::Schema.define do
+        create_table :stagehand_commit_entries do |t|
+          t.integer :record_id
+          t.string :table_name
+          t.string :operation, :null => false
+          t.integer :commit_id
+          t.boolean :committed, :null => false, :default => false
+          t.datetime :created_at
+        end
+
+        add_index :stagehand_commit_entries, :commit_id # Used for looking up all entries within a commit
+        add_index :stagehand_commit_entries, [:record_id, :table_name, :committed], :name => 'index_stagehand_commit_entries_for_matching' # Used for 'matching' scope
+        add_index :stagehand_commit_entries, [:operation, :committed, :commit_id], :name => 'index_stagehand_commit_entries_for_loading' # Used for looking up start entries
+      end
+
+      Stagehand::Staging::CommitEntry.reset_column_information
+
+      add_stagehand!(table_options)
+    end
+
+
+    def add_stagehand!(force: false, **table_options)
+      return if Database.connected_to_production? && !Stagehand::Configuration.single_connection?
+
+      ActiveRecord::Schema.define do
+        Stagehand::Schema.send :each_table, table_options do |table_name|
+          Stagehand::Schema.send :create_operation_trigger, table_name, 'insert', 'NEW', force: force
+          Stagehand::Schema.send :create_operation_trigger, table_name, 'update', 'NEW', force: force
+          Stagehand::Schema.send :create_operation_trigger, table_name, 'delete', 'OLD', force: force
+        end
+      end
+    end
+
+    def remove_stagehand!(remove_entries: true, **table_options)
+      ActiveRecord::Schema.define do
+        Stagehand::Schema.send :each_table, table_options do |table_name|
+          next unless Stagehand::Schema.send :has_stagehand_triggers?, table_name
+          Stagehand::Schema.send :drop_trigger, table_name, 'insert'
+          Stagehand::Schema.send :drop_trigger, table_name, 'update'
+          Stagehand::Schema.send :drop_trigger, table_name, 'delete'
+          Stagehand::Schema.send :expunge, table_name if remove_entries
+        end
+
+        drop_table :stagehand_commit_entries unless table_options[:only].present?
+      end
+    end
+
+    def has_stagehand?(table_name = nil)
+      if UNTRACKED_TABLES.include?(table_name.to_s)
+        return false
+      elsif table_name
+        has_stagehand_triggers?(table_name)
+      else
+        ActiveRecord::Base.connection.table_exists?(Stagehand::Staging::CommitEntry.table_name)
+      end
+    end
+
+    private
+
+    def each_table(only: nil, except: nil)
+      table_names = ActiveRecord::Base.connection.tables
+      table_names -= UNTRACKED_TABLES
+      table_names -= Array(except).collect(&:to_s)
+      table_names &= Array(only).collect(&:to_s) if only.present?
+
+      table_names.each do |table_name|
+        yield table_name
+      end
+    end
+
+    def create_operation_trigger(table_name, trigger_event, record, force: false)
+      if force
+        drop_trigger(table_name, trigger_event)
+      elsif trigger_exists?(table_name, trigger_event)
+        return
+      end
+
+      create_trigger(table_name, trigger_event, :after, <<-SQL)
+        BEGIN
+          INSERT INTO stagehand_commit_entries (record_id, table_name, operation, commit_id, created_at)
+          VALUES (#{record}.id, '#{table_name}', '#{trigger_event}', @stagehand_commit_id, CURRENT_TIMESTAMP());
+        END;
+      SQL
+    end
+
+    def create_trigger(table_name, trigger_event, trigger_time, row_action)
+      ActiveRecord::Base.connection.execute <<-SQL
+        CREATE TRIGGER #{trigger_name(table_name, trigger_event)} #{trigger_time} #{trigger_event}
+        ON #{table_name} FOR EACH ROW #{row_action}
+      SQL
+    end
+
+    def drop_trigger(table_name, trigger_event)
+      ActiveRecord::Base.connection.execute("DROP TRIGGER IF EXISTS #{trigger_name(table_name, trigger_event)};")
+    end
+
+    def trigger_exists?(table_name, trigger_event)
+      ActiveRecord::Base.connection.select_one("SHOW TRIGGERS where `trigger` = '#{trigger_name(table_name, trigger_event)}'").present?
+    end
+
+    def has_stagehand_triggers?(table_name)
+      get_triggers(table_name).present?
+    end
+
+    def trigger_name(table_name, trigger_event)
+      "stagehand_#{trigger_event}_trigger_#{table_name}".downcase
+    end
+
+    def get_triggers(table_name = nil)
+      statement = <<~SQL
+        SHOW TRIGGERS WHERE `Trigger` LIKE 'stagehand_%'
+      SQL
+      statement << " AND `Table` LIKE #{ActiveRecord::Base.connection.quote(table_name)}" if table_name.present?
+
+      return ActiveRecord::Base.connection.select_all(statement)
+    end
+
+    def expunge(table_name)
+      commit_ids = [] # Keep track of commits that we need to clean up if they're now empty
+
+      # Remove records from the table as the subject of any commits
+      Stagehand::Staging::CommitEntry.start_operations.where(:table_name => table_name).in_batches do |batch|
+        commit_ids.concat batch.contained.distinct.pluck(:commit_id)
+        batch.update_all(:record_id => nil, :table_name => nil)
+      end
+
+      # Remove commit entries for records from the table
+      Stagehand::Staging::CommitEntry.content_operations.where(:table_name => table_name).in_batches do |batch|
+        commit_ids.concat batch.contained.distinct.pluck(:commit_id)
+        batch.delete_all
+      end
+
+      Stagehand::Staging::Commit.find(commit_ids).select(&:empty?).each(&:destroy)
+    end
+  end
+end

data/lib/stagehand/schema_extensions.rb

@@ -0,0 +1,10 @@
+module Stagehand
+  module SchemaExtensions
+    def define(*)
+      # Allow production writes during Schema.define to allow Rails to write to ar_internal_metadata table
+      Stagehand::Connection.with_production_writes { super }
+    end
+  end
+end
+
+ActiveRecord::Schema.prepend(Stagehand::SchemaExtensions)