pg_online_schema_change 0.1.0

data/lib/pg_online_schema_change/query.rb

@@ -0,0 +1,302 @@
+require "pg_query"
+require "pg"
+
+module PgOnlineSchemaChange
+  class Query
+    extend Helper
+
+    INDEX_SUFFIX = "_pgosc".freeze
+    DROPPED_COLUMN_TYPE = :AT_DropColumn
+    RENAMED_COLUMN_TYPE = :AT_RenameColumn
+    LOCK_ATTEMPT = 4
+
+    class << self
+      def alter_statement?(query)
+        PgQuery.parse(query).tree.stmts.all? do |statement|
+          statement.stmt.alter_table_stmt.instance_of?(PgQuery::AlterTableStmt) || statement.stmt.rename_stmt.instance_of?(PgQuery::RenameStmt)
+        end
+      rescue PgQuery::ParseError => e
+        false
+      end
+
+      def same_table?(query)
+        tables = PgQuery.parse(query).tree.stmts.map do |statement|
+          if statement.stmt.alter_table_stmt.instance_of?(PgQuery::AlterTableStmt)
+            statement.stmt.alter_table_stmt.relation.relname
+          elsif statement.stmt.rename_stmt.instance_of?(PgQuery::RenameStmt)
+            statement.stmt.rename_stmt.relation.relname
+          end
+        end.compact
+
+        tables.uniq.count == 1
+      rescue PgQuery::ParseError => e
+        false
+      end
+
+      def table(query)
+        from_rename_statement = PgQuery.parse(query).tree.stmts.map do |statement|
+          statement.stmt.rename_stmt&.relation&.relname
+        end.compact[0]
+        PgQuery.parse(query).tables[0] || from_rename_statement
+      end
+
+      def run(connection, query, reuse_trasaction = false, &block)
+        connection.cancel if [PG::PQTRANS_INERROR, PG::PQTRANS_UNKNOWN].include?(connection.transaction_status)
+
+        logger.debug("Running query", { query: query })
+
+        connection.async_exec("BEGIN;")
+
+        result = connection.async_exec(query, &block)
+      rescue Exception
+        connection.cancel if connection.transaction_status != PG::PQTRANS_IDLE
+        connection.block
+        logger.info("Exception raised, rolling back query", { rollback: true, query: query })
+        connection.async_exec("ROLLBACK;")
+        connection.async_exec("COMMIT;")
+        raise
+      else
+        connection.async_exec("COMMIT;") unless reuse_trasaction
+        result
+      end
+
+      def table_columns(client, table = nil)
+        sql = <<~SQL
+          SELECT attname as column_name, format_type(atttypid, atttypmod) as type, attnum as column_position FROM   pg_attribute
+          WHERE  attrelid = \'#{table || client.table}\'::regclass AND attnum > 0 AND NOT attisdropped
+          ORDER  BY attnum;
+        SQL
+        mapped_columns = []
+
+        run(client.connection, sql) do |result|
+          mapped_columns = result.map do |row|
+            row["column_name_regular"] = row["column_name"]
+            row["column_name"] = client.connection.quote_ident(row["column_name"])
+            row["column_position"] = row["column_position"].to_i
+            row
+          end
+        end
+
+        mapped_columns
+      end
+
+      def alter_statement_for(client, shadow_table)
+        parsed_query = PgQuery.parse(client.alter_statement)
+
+        parsed_query.tree.stmts.each do |statement|
+          statement.stmt.alter_table_stmt.relation.relname = shadow_table if statement.stmt.alter_table_stmt
+
+          statement.stmt.rename_stmt.relation.relname = shadow_table if statement.stmt.rename_stmt
+        end
+        parsed_query.deparse
+      end
+
+      def get_indexes_for(client, table)
+        query = <<~SQL
+          SELECT indexdef, schemaname
+          FROM pg_indexes
+          WHERE schemaname = \'#{client.schema}\' AND tablename = \'#{table}\'
+        SQL
+
+        indexes = []
+        run(client.connection, query) do |result|
+          indexes = result.map { |row| row["indexdef"] }
+        end
+
+        indexes
+      end
+
+      def get_all_constraints_for(client)
+        query = <<~SQL
+          SELECT  conrelid::regclass AS table_on,
+                  confrelid::regclass AS table_from,
+                  contype as constraint_type,
+                  conname AS constraint_name,
+                  convalidated AS constraint_validated,
+                  pg_get_constraintdef(oid) AS definition
+          FROM   	pg_constraint
+          WHERE  	contype IN ('f', 'p')
+        SQL
+
+        constraints = []
+        run(client.connection, query) do |result|
+          constraints = result.map { |row| row }
+        end
+
+        constraints
+      end
+
+      def get_primary_keys_for(client, table)
+        get_all_constraints_for(client).select do |row|
+          row["table_on"] == table && row["constraint_type"] == "p"
+        end
+      end
+
+      def get_foreign_keys_for(client, table)
+        get_all_constraints_for(client).select do |row|
+          row["table_on"] == table && row["constraint_type"] == "f"
+        end
+      end
+
+      def get_foreign_keys_to_refresh(client, table)
+        references = get_all_constraints_for(client).select do |row|
+          row["table_from"] == table && row["constraint_type"] == "f"
+        end
+
+        references.map do |row|
+          if row["definition"].end_with?("NOT VALID")
+            add_statement = "ALTER TABLE #{row["table_on"]} ADD CONSTRAINT #{row["constraint_name"]} #{row["definition"]};"
+          else
+            add_statement = "ALTER TABLE #{row["table_on"]} ADD CONSTRAINT #{row["constraint_name"]} #{row["definition"]} NOT VALID;"
+          end
+
+          drop_statement = "ALTER TABLE #{row["table_on"]} DROP CONSTRAINT #{row["constraint_name"]};"
+
+          "#{drop_statement} #{add_statement}"
+        end.join
+      end
+
+      def get_foreign_keys_to_validate(client, table)
+        references = get_all_constraints_for(client).select do |row|
+          row["table_from"] == table && row["constraint_type"] == "f"
+        end
+
+        references.map do |row|
+          "ALTER TABLE #{row["table_on"]} VALIDATE CONSTRAINT #{row["constraint_name"]};"
+        end.join
+      end
+
+      def dropped_columns(client)
+        PgQuery.parse(client.alter_statement).tree.stmts.map do |statement|
+          next if statement.stmt.alter_table_stmt.nil?
+
+          statement.stmt.alter_table_stmt.cmds.map do |cmd|
+            cmd.alter_table_cmd.name if cmd.alter_table_cmd.subtype == DROPPED_COLUMN_TYPE
+          end
+        end.flatten.compact
+      end
+
+      def renamed_columns(client)
+        PgQuery.parse(client.alter_statement).tree.stmts.map do |statement|
+          next if statement.stmt.rename_stmt.nil?
+
+          {
+            old_name: statement.stmt.rename_stmt.subname,
+            new_name: statement.stmt.rename_stmt.newname,
+          }
+        end.flatten.compact
+      end
+
+      def primary_key_for(client, table)
+        query = <<~SQL
+          SELECT
+            pg_attribute.attname as column_name
+          FROM pg_index, pg_class, pg_attribute, pg_namespace
+          WHERE
+            pg_class.oid = \'#{table}\'::regclass AND
+            indrelid = pg_class.oid AND
+            nspname = \'#{client.schema}\' AND
+            pg_class.relnamespace = pg_namespace.oid AND
+            pg_attribute.attrelid = pg_class.oid AND
+            pg_attribute.attnum = any(pg_index.indkey)
+          AND indisprimary
+        SQL
+
+        columns = []
+        run(client.connection, query) do |result|
+          columns = result.map { |row| row["column_name"] }
+        end
+
+        columns.first
+      end
+
+      def storage_parameters_for(client, table)
+        query = <<~SQL
+          SELECT array_to_string(reloptions, ',') as params FROM pg_class WHERE relname=\'#{table}\';
+        SQL
+
+        columns = []
+        run(client.connection, query) do |result|
+          columns = result.map { |row| row["params"] }
+        end
+
+        columns.first
+      end
+
+      # This function acquires the lock and keeps the transaction
+      # open. If a lock is acquired, its upon the caller
+      # to call COMMIT to end the transaction. If a lock
+      # is not acquired, transaction is closed and a new transaction
+      # is started to acquire lock again
+      def open_lock_exclusive(client, table)
+        attempts ||= 1
+
+        query = <<~SQL
+          SET lock_timeout = '#{client.wait_time_for_lock}s';
+          LOCK TABLE #{client.table} IN ACCESS EXCLUSIVE MODE;
+        SQL
+        run(client.connection, query, true)
+
+        true
+      rescue PG::LockNotAvailable, PG::InFailedSqlTransaction
+        if (attempts += 1) < LOCK_ATTEMPT
+          logger.info("Couldn't acquire lock, attempt: #{attempts}")
+
+          run(client.connection, "RESET lock_timeout;")
+          kill_backends(client, table)
+
+          retry
+        end
+
+        logger.info("Lock acquire failed")
+        run(client.connection, "RESET lock_timeout;")
+
+        false
+      end
+
+      def kill_backends(client, table)
+        return unless client.kill_backends
+
+        logger.info("Terminating other backends")
+
+        query = <<~SQL
+          SELECT pg_terminate_backend(pid) FROM pg_locks WHERE locktype = 'relation' AND relation = \'#{table}\'::regclass::oid AND pid <> pg_backend_pid()
+        SQL
+
+        run(client.connection, query, true)
+      end
+
+      def copy_data_statement(client, shadow_table)
+        select_columns = table_columns(client).map do |entry|
+          entry["column_name_regular"]
+        end
+
+        select_columns -= dropped_columns_list if dropped_columns_list.any?
+
+        insert_into_columns = select_columns.dup
+
+        if renamed_columns_list.any?
+          renamed_columns_list.each do |obj|
+            insert_into_columns.each_with_index do |insert_into_column, index|
+              insert_into_columns[index] = obj[:new_name] if insert_into_column == obj[:old_name]
+            end
+          end
+        end
+
+        insert_into_columns.map! do |insert_into_column|
+          client.connection.quote_ident(insert_into_column)
+        end
+
+        select_columns.map! do |select_column|
+          client.connection.quote_ident(select_column)
+        end
+
+        sql = <<~SQL
+          INSERT INTO #{shadow_table}(#{insert_into_columns.join(", ")})
+          SELECT #{select_columns.join(", ")}
+          FROM ONLY #{client.table}
+        SQL
+      end
+    end
+  end
+end

data/lib/pg_online_schema_change/replay.rb

@@ -0,0 +1,122 @@
+module PgOnlineSchemaChange
+  class Replay
+    extend Helper
+
+    class << self
+      PULL_BATCH_COUNT = 1000
+      DELTA_COUNT = 20
+      RESERVED_COLUMNS = %w[operation_type trigger_time].freeze
+
+      # This, picks PULL_BATCH_COUNT rows by primary key from audit_table,
+      # replays it on the shadow_table. Once the batch is done,
+      # it them deletes those PULL_BATCH_COUNT rows from audit_table. Then, pull another batch,
+      # check if the row count matches PULL_BATCH_COUNT, if so swap, otherwise
+      # continue. Swap because, the row count is minimal to replay them altogether
+      # and perform the rename while holding an access exclusive lock for minimal time.
+      def begin!
+        loop do
+          rows = rows_to_play
+
+          raise CountBelowDelta if rows.count <= DELTA_COUNT
+
+          play!(rows)
+        end
+      end
+
+      def rows_to_play(reuse_trasaction = false)
+        select_query = <<~SQL
+          SELECT * FROM #{audit_table} ORDER BY #{primary_key} LIMIT #{PULL_BATCH_COUNT};
+        SQL
+
+        rows = []
+        Query.run(client.connection, select_query, reuse_trasaction) { |result| rows = result.map { |row| row } }
+
+        rows
+      end
+
+      def play!(rows, reuse_trasaction = false)
+        logger.info("Replaying rows, count: #{rows.size}")
+
+        to_be_deleted_rows = []
+        to_be_replayed = []
+        rows.each do |row|
+          new_row = row.dup
+
+          # Remove audit table cols, since we will be
+          # re-mapping them for inserts and updates
+          RESERVED_COLUMNS.each do |col|
+            new_row.delete(col)
+          end
+
+          if dropped_columns_list.any?
+            dropped_columns_list.each do |dropped_column|
+              new_row.delete(dropped_column)
+            end
+          end
+
+          if renamed_columns_list.any?
+            renamed_columns_list.each do |object|
+              value = new_row.delete(object[:old_name])
+              new_row[object[:new_name]] = value
+            end
+          end
+
+          new_row = new_row.compact
+
+          # quote indent column to preserve case insensitivity
+          # ensure rows are escaped
+          new_row = new_row.transform_keys do |column|
+            client.connection.quote_ident(column)
+          end
+
+          new_row = new_row.transform_values do |value|
+            client.connection.escape_string(value)
+          end
+
+          case row["operation_type"]
+          when "INSERT"
+            values = new_row.map { |_, val| "'#{val}'" }.join(",")
+
+            sql = <<~SQL
+              INSERT INTO #{shadow_table} (#{new_row.keys.join(",")})
+                VALUES (#{values});
+            SQL
+            to_be_replayed << sql
+
+            to_be_deleted_rows << "'#{row[primary_key]}'"
+          when "UPDATE"
+            set_values = new_row.map do |column, value|
+              "#{column} = '#{value}'"
+            end.join(",")
+
+            sql = <<~SQL
+              UPDATE #{shadow_table}
+              SET #{set_values}
+              WHERE #{primary_key}=\'#{row[primary_key]}\';
+            SQL
+            to_be_replayed << sql
+
+            to_be_deleted_rows << "'#{row[primary_key]}'"
+          when "DELETE"
+            sql = <<~SQL
+              DELETE FROM #{shadow_table} WHERE #{primary_key}=\'#{row[primary_key]}\';
+            SQL
+            to_be_replayed << sql
+
+            to_be_deleted_rows << "'#{row[primary_key]}'"
+          end
+        end
+
+        Query.run(client.connection, to_be_replayed.join, reuse_trasaction)
+
+        # Delete items from the audit now that are replayed
+        if to_be_deleted_rows.count >= 1
+          delete_query = <<~SQL
+            DELETE FROM #{audit_table} WHERE #{primary_key} IN (#{to_be_deleted_rows.join(",")})
+          SQL
+          Query.run(client.connection, delete_query, reuse_trasaction)
+        end
+      end
+    end
+  end
+end

data/lib/pg_online_schema_change/store.rb

@@ -0,0 +1,18 @@
+require "pg_query"
+require "pg"
+
+module PgOnlineSchemaChange
+  class Store
+    class << self
+      @@object = {}
+
+      def get(key)
+        @@object[key.to_s] || @@object[key.to_sym]
+      end
+
+      def set(key, value)
+        @@object[key.to_sym] = value
+      end
+    end
+  end
+end

data/lib/pg_online_schema_change/version.rb

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

data/lib/pg_online_schema_change.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+require "json"
+require "ougai"
+
+require "pg_online_schema_change/version"
+require "pg_online_schema_change/helper"
+require "pg_online_schema_change/functions"
+require "pg_online_schema_change/cli"
+require "pg_online_schema_change/client"
+require "pg_online_schema_change/query"
+require "pg_online_schema_change/store"
+require "pg_online_schema_change/replay"
+require "pg_online_schema_change/orchestrate"
+
+module PgOnlineSchemaChange
+  class Error < StandardError; end
+  class CountBelowDelta < StandardError; end
+  class AccessExclusiveLockNotAcquired < StandardError; end
+
+  def self.logger=(verbose)
+    @@logger ||= begin
+        logger = Ougai::Logger.new($stdout)
+        logger.level = verbose ? Ougai::Logger::TRACE : Ougai::Logger::INFO
+        logger.with_fields = { version: PgOnlineSchemaChange::VERSION }
+        logger
+      end
+  end
+
+  def self.logger
+    @@logger
+  end
+end

metadata

@@ -0,0 +1,76 @@
+--- !ruby/object:Gem::Specification
+name: pg_online_schema_change
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- Shayon Mukherjee
+autorequire:
+bindir: bin
+cert_chain: []
+date: 2022-02-16 00:00:00.000000000 Z
+dependencies: []
+description: pg-online-schema-change (pg-osc) is a tool for making schema changes
+  (any ALTER statements) in Postgres tables with minimal locks, thus helping achieve
+  zero downtime schema changes against production workloads.
+email:
+- shayonj@gmail.com
+executables:
+- console
+- pg-online-schema-change
+- setup
+extensions: []
+extra_rdoc_files: []
+files:
+- ".rspec"
+- ".rubocop.yml"
+- ".ruby-version"
+- CHANGELOG.md
+- CODE_OF_CONDUCT.md
+- Gemfile
+- Gemfile.lock
+- LICENSE.txt
+- README.md
+- Rakefile
+- bin/console
+- bin/pg-online-schema-change
+- bin/setup
+- diagrams/how-it-works.excalidraw
+- diagrams/how-it-works.png
+- docker-compose.yml
+- lib/pg_online_schema_change.rb
+- lib/pg_online_schema_change/cli.rb
+- lib/pg_online_schema_change/client.rb
+- lib/pg_online_schema_change/functions.rb
+- lib/pg_online_schema_change/helper.rb
+- lib/pg_online_schema_change/orchestrate.rb
+- lib/pg_online_schema_change/query.rb
+- lib/pg_online_schema_change/replay.rb
+- lib/pg_online_schema_change/store.rb
+- lib/pg_online_schema_change/version.rb
+homepage: https://github.com/shayonj/pg-online-schema-change
+licenses:
+- MIT
+metadata:
+  rubygems_mfa_required: 'true'
+post_install_message:
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: 2.6.0
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.2.3
+signing_key:
+specification_version: 4
+summary: pg-online-schema-change is a tool for schema changes for Postgres tables
+  with minimal locks
+test_files: []