pg_online_schema_change 0.1.0

data/docker-compose.yml

@@ -0,0 +1,9 @@
+version: '2'
+services:
+  postgres:
+    image: postgres:9.6.2-alpine
+    ports:
+      - "5432:5432"
+    environment:
+      POSTGRES_USER: jamesbond
+      POSTGRES_DB: postgres

data/lib/pg_online_schema_change/cli.rb

@@ -0,0 +1,44 @@
+require "thor"
+
+module PgOnlineSchemaChange
+  class CLI < Thor
+    desc "perform", "Perform the set of operations to safely apply the schema change with minimal locks"
+    method_option :alter_statement, aliases: "-a", type: :string, required: true,
+                                    desc: "The ALTER statement to perform the schema change"
+    method_option :schema, aliases: "-s", type: :string, required: true, default: "public",
+                           desc: "The schema in which the table is"
+    method_option :dbname, aliases: "-d", type: :string, required: true, desc: "Name of the database"
+    method_option :host, aliases: "-h", type: :string, required: true, desc: "Server host where the Database is located"
+    method_option :username, aliases: "-u", type: :string, required: true, desc: "Username for the Database"
+    method_option :port, aliases: "-p", type: :numeric, required: true, default: 5432, desc: "Port for the Database"
+    method_option :password, aliases: "-w", type: :string, required: true, desc: "Password for the Database"
+    method_option :verbose, aliases: "-v", type: :boolean, default: false, desc: "Emit logs in debug mode"
+    method_option :drop, aliases: "-f", type: :boolean, default: false,
+                         desc: "Drop the original table in the end after the swap"
+    method_option :kill_backends, aliases: "-k", type: :boolean, default: false,
+                                  desc: "Kill other competing queries/backends when trying to acquire lock for the shadow table creation and swap. It will wait for --wait-time-for-lock duration before killing backends and try upto 3 times."
+    method_option :wait_time_for_lock, aliases: "-w", type: :numeric, default: 10,
+                                       desc: "Time to wait before killing backends to acquire lock and/or retrying upto 3 times. It will kill backends if --kill-backends is true, otherwise try upto 3 times and exit if it cannot acquire a lock."
+    method_option :copy_statement, aliases: "-c", type: :string, required: false, default: "",
+                                   desc: "Takes a .sql file location where you can provide a custom query to be played (ex: backfills) when pgosc copies data from the primary to the shadow table. More examples in README."
+
+    def perform
+      client_options = Struct.new(*options.keys.map(&:to_sym)).new(*options.values)
+
+      PgOnlineSchemaChange.logger = client_options.verbose
+      PgOnlineSchemaChange::Orchestrate.run!(client_options)
+    end
+
+    desc "--version, -v", "print the version"
+
+    def version
+      puts PgOnlineSchemaChange::VERSION
+    end
+
+    def self.exit_on_failure?
+      true
+    end
+
+    default_task :perform
+  end
+end

data/lib/pg_online_schema_change/client.rb

@@ -0,0 +1,49 @@
+require "pg"
+
+module PgOnlineSchemaChange
+  class Client
+    attr_accessor :alter_statement, :schema, :dbname, :host, :username, :port, :password, :connection, :table, :drop,
+                  :kill_backends, :wait_time_for_lock, :copy_statement
+
+    def initialize(options)
+      @alter_statement = options.alter_statement
+      @schema = options.schema
+      @dbname = options.dbname
+      @host = options.host
+      @username = options.username
+      @port = options.port
+      @password = options.password
+      @drop = options.drop
+      @kill_backends = options.kill_backends
+      @wait_time_for_lock = options.wait_time_for_lock
+      handle_copy_statement(options.copy_statement)
+
+      @connection = PG.connect(
+        dbname: @dbname,
+        host: @host,
+        user: @username,
+        password: @password,
+        port: @port,
+      )
+
+      raise Error, "Not a valid ALTER statement: #{@alter_statement}" unless Query.alter_statement?(@alter_statement)
+
+      unless Query.same_table?(@alter_statement)
+        raise Error "All statements should belong to the same table: #{@alter_statement}"
+      end
+
+      @table = Query.table(@alter_statement)
+
+      PgOnlineSchemaChange.logger.debug("Connection established")
+    end
+
+    def handle_copy_statement(statement)
+      return if statement.nil? || statement == ""
+
+      file_path = File.expand_path(statement)
+      raise Error, "File not found: #{file_path}" unless File.file?(file_path)
+
+      @copy_statement = File.open(file_path, "rb", &:read)
+    end
+  end
+end

data/lib/pg_online_schema_change/functions.rb

@@ -0,0 +1,62 @@
+FUNC_FIX_SERIAL_SEQUENCE = <<~SQL.freeze
+    CREATE OR REPLACE FUNCTION fix_serial_sequence(_table regclass, _newtable text)
+    RETURNS void AS
+    $func$
+    DECLARE
+    _sql text;
+    BEGIN
+
+    -- Update serial columns to ensure copied table doesn't follow same sequence as primary table
+    SELECT INTO _sql
+        string_agg('CREATE SEQUENCE ' || seq, E';\n') || E';\n'
+        || string_agg(format('ALTER SEQUENCE %s OWNED BY %I.%I'
+                            , seq, _newtable, a.attname), E';\n') || E';\n'
+        || 'ALTER TABLE ' || quote_ident(_newtable) || E'\n  '
+        || string_agg(format($$ALTER %I SET DEFAULT nextval('%s'::regclass)$$
+                                    , a.attname, seq), E'\n, ')
+    FROM   pg_attribute  a
+    JOIN   pg_attrdef    ad ON ad.adrelid = a.attrelid
+                        AND ad.adnum   = a.attnum
+        , quote_ident(_newtable || '_' || a.attname || '_seq') AS seq
+    WHERE  a.attrelid = _table
+    AND    a.attnum > 0
+    AND    NOT a.attisdropped
+    AND    a.atttypid = ANY ('{int,int8,int2}'::regtype[])
+    AND    pg_get_expr(ad.adbin, ad.adrelid) = 'nextval('''
+          || (pg_get_serial_sequence (a.attrelid::regclass::text, a.attname))::regclass
+          || '''::regclass)'
+  ;
+
+  IF _sql IS NOT NULL THEN
+  EXECUTE _sql;
+  END IF;
+
+  END
+  $func$  LANGUAGE plpgsql VOLATILE;
+SQL
+
+FUNC_CREATE_TABLE_ALL = <<~SQL.freeze
+  CREATE OR REPLACE FUNCTION create_table_all(source_table text, newsource_table text)
+    RETURNS void language plpgsql
+    as $$
+    declare
+        rec record;
+    begin
+    EXECUTE format(
+          'CREATE TABLE %s (LIKE %s including all)',
+          newsource_table, source_table);
+      for rec in
+          SELECT oid, conname
+          FROM pg_constraint
+          WHERE contype = 'f'
+          AND conrelid = source_table::regclass
+      LOOP
+      EXECUTE format(
+              'ALTER TABLE %s add constraint %s %s',
+              newsource_table,
+              rec.conname,
+              pg_get_constraintdef(rec.oid));
+      END LOOP;
+    END
+  $$;
+SQL

data/lib/pg_online_schema_change/helper.rb

@@ -0,0 +1,21 @@
+module PgOnlineSchemaChange
+  module Helper
+    def primary_key
+      result = Store.get(:primary_key)
+      return result if result
+
+      Store.set(:primary_key, Query.primary_key_for(client, client.table))
+    end
+
+    def logger
+      PgOnlineSchemaChange.logger
+    end
+
+    def method_missing(method, *_args)
+      result = Store.send(:get, method)
+      return result if result
+
+      raise ArgumentError, "Method `#{method}` doesn't exist."
+    end
+  end
+end

data/lib/pg_online_schema_change/orchestrate.rb

@@ -0,0 +1,263 @@
+require "securerandom"
+
+module PgOnlineSchemaChange
+  class Orchestrate
+    SWAP_STATEMENT_TIMEOUT = "5s".freeze
+
+    extend Helper
+
+    class << self
+      def setup!(options)
+        client = Store.set(:client, Client.new(options))
+
+        sql = <<~SQL
+          SET statement_timeout = 0;
+          SET client_min_messages = warning;
+          SET search_path TO #{client.schema};
+        SQL
+
+        Query.run(client.connection, sql)
+        # install functions
+        Query.run(client.connection, FUNC_FIX_SERIAL_SEQUENCE)
+        Query.run(client.connection, FUNC_CREATE_TABLE_ALL)
+
+        # Set this early on to ensure their creation and cleanup (unexpected)
+        # happens at all times. IOW, the calls from Store.get always return
+        # the same value.
+        Store.set(:old_primary_table, "pgosc_op_table_#{client.table}")
+        Store.set(:audit_table, "pgosc_at_#{client.table}_#{random_string}")
+        Store.set(:shadow_table, "pgosc_st_#{client.table}_#{random_string}")
+      end
+
+      def run!(options)
+        setup!(options)
+        Thread.new { handle_signals! }
+
+        raise Error, "Parent table has no primary key, exiting..." if primary_key.nil?
+
+        setup_audit_table!
+        setup_trigger!
+        setup_shadow_table!
+        disable_vacuum!
+        run_alter_statement!
+        copy_data!
+        run_analyze!
+        replay_and_swap!
+        run_analyze!
+        validate_constraints!
+        drop_and_cleanup!
+
+        logger.info("All tasks successfully completed")
+      rescue StandardError => e
+        logger.fatal("Something went wrong: #{e.message}", { e: e })
+
+        drop_and_cleanup!
+
+        raise e
+      end
+
+      def setup_signals!
+        reader, writer = IO.pipe
+
+        %w[TERM QUIT INT].each do |sig|
+          trap(sig) { writer.puts sig }
+        end
+
+        reader
+      end
+
+      def handle_signals!
+        reader = setup_signals!
+        signal = reader.gets.chomp
+
+        while !reader.closed? && IO.select([reader])
+          logger.info "Signal #{signal} received, cleaning up"
+
+          client.connection.cancel
+          drop_and_cleanup!
+          reader.close
+
+          exit Signal.list[signal]
+        end
+      end
+
+      def setup_audit_table!
+        logger.info("Setting up audit table", { audit_table: audit_table })
+
+        sql = <<~SQL
+          CREATE TABLE #{audit_table} (operation_type text, trigger_time timestamp, LIKE #{client.table});
+        SQL
+
+        Query.run(client.connection, sql)
+      end
+
+      def setup_trigger!
+        # acquire access exclusive lock to ensure audit triggers
+        # are setup fine. This also calls kill_backends (if opted in via flag)
+        # so any competing backends will be killed to setup the trigger
+        opened = Query.open_lock_exclusive(client, client.table)
+
+        raise AccessExclusiveLockNotAcquired unless opened
+
+        logger.info("Setting up triggers")
+
+        sql = <<~SQL
+          DROP TRIGGER IF EXISTS primary_to_audit_table_trigger ON #{client.table};
+
+          CREATE OR REPLACE FUNCTION primary_to_audit_table_trigger()
+          RETURNS TRIGGER AS
+          $$
+          BEGIN
+            IF ( TG_OP = 'INSERT') THEN
+              INSERT INTO \"#{audit_table}\" select 'INSERT', now(), NEW.* ;
+              RETURN NEW;
+            ELSIF ( TG_OP = 'UPDATE') THEN
+              INSERT INTO \"#{audit_table}\" select 'UPDATE', now(),  NEW.* ;
+              RETURN NEW;
+            ELSIF ( TG_OP = 'DELETE') THEN
+              INSERT INTO \"#{audit_table}\" select 'DELETE', now(), OLD.* ;
+              RETURN NEW;
+            END IF;
+          END;
+          $$ LANGUAGE PLPGSQL SECURITY DEFINER;
+
+          CREATE TRIGGER primary_to_audit_table_trigger
+          AFTER INSERT OR UPDATE OR DELETE ON #{client.table}
+          FOR EACH ROW EXECUTE PROCEDURE primary_to_audit_table_trigger();
+        SQL
+
+        Query.run(client.connection, sql, opened)
+      ensure
+        Query.run(client.connection, "COMMIT;")
+      end
+
+      def setup_shadow_table!
+        logger.info("Setting up shadow table", { shadow_table: shadow_table })
+
+        Query.run(client.connection, "SELECT create_table_all('#{client.table}', '#{shadow_table}');")
+
+        # update serials
+        Query.run(client.connection, "SELECT fix_serial_sequence('#{client.table}', '#{shadow_table}');")
+      end
+
+      # Disabling vacuum to avoid any issues during the process
+      def disable_vacuum!
+        result = Query.storage_parameters_for(client, client.table) || ""
+        primary_table_storage_parameters = Store.set(:primary_table_storage_parameters, result)
+
+        logger.debug("Disabling vacuum on shadow and audit table",
+                     { shadow_table: shadow_table, audit_table: audit_table })
+        sql = <<~SQL
+          ALTER TABLE #{shadow_table} SET (
+            autovacuum_enabled = false, toast.autovacuum_enabled = false
+          );
+
+          ALTER TABLE #{audit_table} SET (
+            autovacuum_enabled = false, toast.autovacuum_enabled = false
+          );
+        SQL
+        Query.run(client.connection, sql)
+      end
+
+      def run_alter_statement!
+        statement = Query.alter_statement_for(client, shadow_table)
+        logger.info("Running alter statement on shadow table",
+                    { shadow_table: shadow_table, parent_table: client.table })
+        Query.run(client.connection, statement)
+
+        Store.set(:dropped_columns_list, Query.dropped_columns(client))
+        Store.set(:renamed_columns_list, Query.renamed_columns(client))
+      end
+
+      # Begin the process to copy data into copy table
+      # depending on the size of the table, this can be a time
+      # taking operation.
+      def copy_data!
+        logger.info("Copying contents..", { shadow_table: shadow_table, parent_table: client.table })
+
+        if client.copy_statement
+          query = format(client.copy_statement, shadow_table: shadow_table)
+          return Query.run(client.connection, query)
+        end
+
+        sql = Query.copy_data_statement(client, shadow_table)
+        Query.run(client.connection, sql)
+      end
+
+      def replay_and_swap!
+        Replay.begin!
+      rescue CountBelowDelta
+        logger.info("Remaining rows below delta count, proceeding towards swap")
+
+        swap!
+      end
+
+      def swap!
+        logger.info("Performing swap!")
+
+        foreign_key_statements = Query.get_foreign_keys_to_refresh(client, client.table)
+        storage_params_reset = primary_table_storage_parameters.empty? ? "" : "ALTER TABLE #{client.table} SET (#{primary_table_storage_parameters});"
+
+        # From here on, all statements are carried out in a single
+        # transaction with access exclusive lock
+
+        opened = Query.open_lock_exclusive(client, client.table)
+
+        raise AccessExclusiveLockNotAcquired unless opened
+
+        Query.run(client.connection, "SET statement_timeout to '#{SWAP_STATEMENT_TIMEOUT}';", opened)
+
+        rows = Replay.rows_to_play(opened)
+        Replay.play!(rows, opened)
+
+        sql = <<~SQL
+          ALTER TABLE #{client.table} RENAME to #{old_primary_table};
+          ALTER TABLE #{shadow_table} RENAME to #{client.table};
+          #{foreign_key_statements}
+          #{storage_params_reset}
+          DROP TRIGGER IF EXISTS primary_to_audit_table_trigger ON #{client.table};
+        SQL
+
+        Query.run(client.connection, sql)
+      ensure
+        Query.run(client.connection, "COMMIT;")
+        Query.run(client.connection, "SET statement_timeout = 0;")
+      end
+
+      def run_analyze!
+        logger.info("Performing ANALYZE!")
+
+        Query.run(client.connection, "ANALYZE VERBOSE #{client.table};")
+      end
+
+      def validate_constraints!
+        logger.info("Validating constraints!")
+
+        validate_statements = Query.get_foreign_keys_to_validate(client, client.table)
+
+        Query.run(client.connection, validate_statements)
+      end
+
+      def drop_and_cleanup!
+        primary_drop = client.drop ? "DROP TABLE IF EXISTS #{old_primary_table};" : ""
+        audit_table_drop = audit_table ? "DROP TABLE IF EXISTS #{audit_table}" : ""
+        shadow_table_drop = shadow_table ? "DROP TABLE IF EXISTS #{shadow_table}" : ""
+
+        sql = <<~SQL
+          #{audit_table_drop};
+          #{shadow_table_drop};
+          #{primary_drop}
+          RESET statement_timeout;
+          RESET client_min_messages;
+          RESET lock_timeout;
+        SQL
+
+        Query.run(client.connection, sql)
+      end
+
+      private def random_string
+        @random_string ||= SecureRandom.hex(3)
+      end
+    end
+  end
+end