pg_easy_replicate 0.1.0

data/lib/pg_easy_replicate/orchestrate.rb

@@ -0,0 +1,308 @@
+# frozen_string_literal: true
+
+module PgEasyReplicate
+  class Orchestrate
+    extend Helper
+
+    class << self
+      DEFAULT_LAG = 200_000 # 200kb
+      DEFAULT_WAIT = 5 # seconds
+
+      def start_sync(options)
+        PgEasyReplicate.assert_config
+
+        create_publication(
+          group_name: options[:group_name],
+          conn_string: source_db_url,
+        )
+
+        add_tables_to_publication(
+          group_name: options[:group_name],
+          tables: options[:tables],
+          conn_string: source_db_url,
+          schema: options[:schema],
+        )
+
+        create_subscription(
+          group_name: options[:group_name],
+          source_conn_string: secondary_source_db_url || source_db_url,
+          target_conn_string: target_db_url,
+        )
+
+        Group.create(
+          name: options[:group_name],
+          table_names: options[:tables],
+          schema_name: options[:schema],
+          started_at: Time.now.utc,
+        )
+      rescue => e
+        stop_sync(
+          group_name: options[:group_name],
+          source_conn_string: source_db_url,
+          target_conn_string: target_db_url,
+        )
+
+        if Group.find(options[:group_name])
+          Group.update(group_name: options[:group_name], failed_at: Time.now)
+        else
+          Group.create(
+            name: options[:group_name],
+            table_names: options[:tables],
+            schema_name: options[:schema],
+            started_at: Time.now.utc,
+            failed_at: Time.now.utc,
+          )
+        end
+
+        abort_with("Starting sync failed: #{e.message}")
+      end
+
+      def create_publication(group_name:, conn_string:)
+        logger.info(
+          "Setting up publication",
+          { publication_name: publication_name(group_name) },
+        )
+        Query.run(
+          query: "create publication #{publication_name(group_name)}",
+          connection_url: conn_string,
+        )
+      end
+
+      def add_tables_to_publication(
+        schema:,
+        group_name:,
+        conn_string:,
+        tables: ""
+      )
+        logger.info(
+          "Adding tables up publication",
+          { publication_name: publication_name(group_name) },
+        )
+        tables = tables&.split(",") || []
+        unless tables.size > 0
+          tables = list_all_tables(schema: schema, conn_string: conn_string)
+        end
+
+        tables.map do |table_name|
+          Query.run(
+            query:
+              "ALTER PUBLICATION #{publication_name(group_name)} ADD TABLE \"#{table_name}\"",
+            connection_url: conn_string,
+            schema: schema,
+          )
+        end
+      end
+
+      def list_all_tables(schema:, conn_string:)
+        Query
+          .run(
+            query:
+              "SELECT table_name FROM information_schema.tables WHERE table_schema = '#{schema}'",
+            connection_url: conn_string,
+          )
+          .map(&:values)
+          .flatten
+      end
+
+      def drop_publication(group_name:, conn_string:)
+        logger.info(
+          "Dropping publication",
+          { publication_name: publication_name(group_name) },
+        )
+        Query.run(
+          query: "DROP PUBLICATION IF EXISTS #{publication_name(group_name)}",
+          connection_url: conn_string,
+        )
+      end
+
+      def create_subscription(
+        group_name:,
+        source_conn_string:,
+        target_conn_string:
+      )
+        logger.info(
+          "Setting up subscription",
+          {
+            publication_name: publication_name(group_name),
+            subscription_name: subscription_name(group_name),
+          },
+        )
+
+        Query.run(
+          query:
+            "CREATE SUBSCRIPTION #{subscription_name(group_name)} CONNECTION '#{source_conn_string}' PUBLICATION #{publication_name(group_name)}",
+          connection_url: target_conn_string,
+          transaction: false,
+        )
+      rescue Sequel::DatabaseError => e
+        if e.message.include?("canceling statement due to statement timeout")
+          abort_with(
+            "Subscription creation failed, please ensure both databases are in the same network region: #{e.message}",
+          )
+        end
+
+        raise
+      end
+
+      def drop_subscription(group_name:, target_conn_string:)
+        logger.info(
+          "Dropping subscription",
+          {
+            publication_name: publication_name(group_name),
+            subscription_name: subscription_name(group_name),
+          },
+        )
+        Query.run(
+          query: "DROP SUBSCRIPTION IF EXISTS #{subscription_name(group_name)}",
+          connection_url: target_conn_string,
+          transaction: false,
+        )
+      end
+
+      def stop_sync(target_conn_string:, source_conn_string:, group_name:)
+        PgEasyReplicate.assert_config
+
+        logger.info(
+          "Stopping sync",
+          {
+            publication_name: publication_name(group_name),
+            subscription_name: subscription_name(group_name),
+          },
+        )
+        drop_publication(
+          group_name: group_name,
+          conn_string: source_conn_string,
+        )
+        drop_subscription(
+          group_name: group_name,
+          target_conn_string: target_conn_string,
+        )
+      end
+
+      def switchover(
+        group_name:,
+        source_conn_string: source_db_url,
+        target_conn_string: target_db_url,
+        lag_delta_size: DEFAULT_LAG
+      )
+        PgEasyReplicate.assert_config
+        group = Group.find(group_name)
+
+        watch_lag(group_name: group_name, lag: lag_delta_size)
+        revoke_connections_on_source_db(group_name)
+        wait_for_remaining_catchup(group_name)
+        refresh_sequences(
+          conn_string: target_conn_string,
+          schema: group[:schema_name],
+        )
+        mark_switchover_complete(group_name)
+        drop_subscription(
+          group_name: group_name,
+          target_conn_string: target_conn_string,
+        )
+      rescue => e
+        restore_connections_on_source_db(group_name)
+
+        abort_with("Switchover sync failed: #{e.message}")
+      end
+
+      def watch_lag(group_name:, wait_time: DEFAULT_WAIT, lag: DEFAULT_LAG)
+        logger.info("Watching lag stats")
+
+        loop do
+          sleep(wait_time)
+
+          unless Stats.all_tables_replicating?(group_name)
+            logger.debug(
+              "All tables haven't reached replicating state, skipping check",
+            )
+            next
+          end
+
+          lag_stat = Stats.lag_stats(group_name).first
+          if lag_stat[:write_lag].nil? || lag_stat[:flush_lag].nil? ||
+               lag_stat[:replay_lag].nil?
+            next
+          end
+
+          logger.debug("Current lag stats: #{lag_stat}")
+
+          below_write_lag = lag_stat[:write_lag] <= lag
+          below_flush_lag = lag_stat[:flush_lag] <= lag
+          below_replay_lag = lag_stat[:replay_lag] <= lag
+
+          break if below_write_lag && below_flush_lag && below_replay_lag
+        end
+
+        logger.info("Lag below #{DEFAULT_LAG} bytes. Continuing...")
+      end
+
+      def wait_for_remaining_catchup(group_name)
+        logger.info("Waiting for remaining WAL to get flushed")
+
+        watch_lag(group_name: group_name, lag: 0, wait_time: 0.2)
+
+        logger.info("Caught up on remaining WAL lag")
+      end
+
+      def revoke_connections_on_source_db(group_name)
+        logger.info(
+          "Lag is now below #{DEFAULT_LAG}, marking source DB to read only",
+        )
+
+        alter_sql =
+          "ALTER USER #{db_user(source_db_url)} set default_transaction_read_only = true"
+        Query.run(query: alter_sql, connection_url: source_db_url)
+
+        kill_sql =
+          "SELECT pg_terminate_backend(pg_stat_activity.pid) FROM pg_stat_activity WHERE usename = '#{db_user(source_db_url)}';"
+
+        Query.run(query: kill_sql, connection_url: source_db_url)
+      end
+
+      def restore_connections_on_source_db(group_name)
+        logger.info("Restoring connections")
+
+        alter_sql =
+          "ALTER USER #{db_user(source_db_url)} set default_transaction_read_only = false"
+        Query.run(query: alter_sql, connection_url: source_db_url)
+      end
+
+      def refresh_sequences(conn_string:, schema: nil)
+        logger.info("Refreshing sequences")
+        sql = <<~SQL
+          DO $$
+          DECLARE
+          i TEXT;
+          BEGIN
+            FOR i IN (
+              SELECT 'SELECT SETVAL('
+                  || quote_literal(quote_ident(PGT.schemaname) || '.' || quote_ident(S.relname))
+                  || ', COALESCE(MAX(' ||quote_ident(C.attname)|| '), 1) ) FROM '
+                  || quote_ident(PGT.schemaname)|| '.'||quote_ident(T.relname)|| ';'
+                FROM pg_class AS S,
+                    pg_depend AS D,
+                    pg_class AS T,
+                    pg_attribute AS C,
+                    pg_tables AS PGT
+              WHERE S.relkind = 'S'
+                AND S.oid = D.objid
+                AND D.refobjid = T.oid
+                AND D.refobjid = C.attrelid
+                AND D.refobjsubid = C.attnum
+                AND T.relname = PGT.tablename
+            ) LOOP
+                EXECUTE i;
+            END LOOP;
+          END $$;
+        SQL
+
+        Query.run(query: sql, connection_url: conn_string, schema: schema)
+      end
+
+      def mark_switchover_complete(group_name)
+        Group.update(group_name: group_name, switchover_completed_at: Time.now)
+      end
+    end
+  end
+end

data/lib/pg_easy_replicate/query.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+module PgEasyReplicate
+  class Query
+    extend Helper
+
+    class << self
+      def run(
+        query:,
+        connection_url:,
+        user: internal_user_name,
+        schema: nil,
+        transaction: true
+      )
+        conn =
+          connect(connection_url: connection_url, schema: schema, user: user)
+        if transaction
+          r =
+            conn.transaction do
+              conn.run("SET search_path to #{schema}") if schema
+              conn.run("SET statement_timeout to '5s'")
+              conn.fetch(query).to_a
+            end
+        else
+          conn.run("SET search_path to #{schema}") if schema
+          conn.run("SET statement_timeout to '5s'")
+          r = conn.fetch(query).to_a
+        end
+        conn.disconnect
+        r
+      ensure
+        conn&.fetch("RESET statement_timeout")
+        conn&.disconnect
+      end
+
+      def connect(connection_url:, user: internal_user_name, schema: nil)
+        c =
+          Sequel.connect(
+            connection_url,
+            user: user,
+            logger: ENV.fetch("DEBUG", nil) ? logger : nil,
+            search_path: schema,
+          )
+        logger.debug("Connection established")
+        c
+      end
+    end
+  end
+end

data/lib/pg_easy_replicate/stats.rb

@@ -0,0 +1,128 @@
+# frozen_string_literal: true
+
+module PgEasyReplicate
+  class Stats
+    REPLICATION_STATE_MAP = {
+      "i" => "initializing",
+      "d" => "data_is_being_copied",
+      "f" => "finished_table_copy",
+      "s" => "synchronized",
+      "r" => "replicating",
+    }.freeze
+    extend Helper
+
+    class << self
+      def object(group_name)
+        PgEasyReplicate.assert_config
+        stats = replication_stats(group_name)
+        group = Group.find(group_name)
+        {
+          lag_stats: lag_stats(group_name),
+          replication_slots: pg_replication_slots(group_name),
+          replication_stats: stats,
+          replication_stats_count_by_state:
+            replication_stats_count_by_state(stats),
+          message_lsn_receipts: message_lsn_receipts(group_name),
+          sync_started_at: group[:started_at],
+          sync_failed_at: group[:failed_at],
+          switchover_completed_at: group[:switchover_completed_at],
+        }
+      end
+
+      def print(group_name)
+        puts JSON.pretty_generate(object(group_name))
+      end
+
+      def follow(group_name)
+        loop do
+          print(group_name)
+          sleep(1)
+        end
+      end
+
+      # Get
+      def lag_stats(group_name)
+        sql = <<~SQL
+          SELECT pid,
+          client_addr,
+          usename as user_name,
+          application_name,
+          state,
+          sync_state,
+          pg_wal_lsn_diff(sent_lsn, write_lsn) AS write_lag,
+          pg_wal_lsn_diff(sent_lsn, flush_lsn) AS flush_lag,
+          pg_wal_lsn_diff(sent_lsn, replay_lsn) AS replay_lag
+          FROM pg_stat_replication
+          WHERE application_name = '#{subscription_name(group_name)}';
+        SQL
+
+        Query.run(query: sql, connection_url: source_db_url)
+      end
+
+      def pg_replication_slots(group_name)
+        sql = <<~SQL
+         select * from pg_replication_slots WHERE slot_name = '#{subscription_name(group_name)}';
+        SQL
+
+        Query.run(query: sql, connection_url: source_db_url)
+      end
+
+      def replication_stats(group_name)
+        sql = <<~SQL
+          SELECT
+          s.subname AS subscription_name,
+          c.relnamespace :: regnamespace :: text as table_schema,
+          c.relname as table_name,
+          rel.srsubstate as replication_state
+        FROM
+          pg_catalog.pg_subscription s
+          JOIN pg_catalog.pg_subscription_rel rel ON rel.srsubid = s.oid
+          JOIN pg_catalog.pg_class c on c.oid = rel.srrelid
+        WHERE s.subname = '#{subscription_name(group_name)}'
+        SQL
+
+        Query
+          .run(query: sql, connection_url: target_db_url)
+          .each do |obj|
+            obj[:replication_state] = REPLICATION_STATE_MAP[
+              obj[:replication_state]
+            ]
+          end
+      end
+
+      def all_tables_replicating?(group_name)
+        result =
+          replication_stats(group_name)
+            .each
+            .with_object(Hash.new(0)) do |state, counts|
+              counts[state[:replication_state]] += 1
+            end
+        result.keys.uniq.count == 1 &&
+          result.keys.first == REPLICATION_STATE_MAP["r"]
+      end
+
+      def replication_stats_count_by_state(stats)
+        stats
+          .each
+          .with_object(Hash.new(0)) do |state, counts|
+            counts[state[:replication_state]] += 1
+          end
+      end
+
+      def message_lsn_receipts(group_name)
+        sql = <<~SQL
+          select
+          received_lsn,
+          last_msg_send_time,
+          last_msg_receipt_time,
+          latest_end_lsn,
+          latest_end_time
+          from
+            pg_catalog.pg_stat_subscription
+          WHERE subname = '#{subscription_name(group_name)}'
+        SQL
+        Query.run(query: sql, connection_url: target_db_url)
+      end
+    end
+  end
+end

data/lib/pg_easy_replicate/version.rb

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

data/lib/pg_easy_replicate.rb

@@ -0,0 +1,197 @@
+# frozen_string_literal: true
+
+require "json"
+require "ougai"
+require "lockbox"
+require "pg"
+require "sequel"
+
+require "pg_easy_replicate/helper"
+require "pg_easy_replicate/version"
+require "pg_easy_replicate/query"
+require "pg_easy_replicate/orchestrate"
+require "pg_easy_replicate/stats"
+require "pg_easy_replicate/group"
+require "pg_easy_replicate/cli"
+
+Sequel.default_timezone = :utc
+module PgEasyReplicate
+  class Error < StandardError
+  end
+
+  extend Helper
+
+  class << self
+    def config
+      abort_with("SOURCE_DB_URL is missing") if source_db_url.nil?
+      abort_with("TARGET_DB_URL is missing") if target_db_url.nil?
+      @config ||=
+        begin
+          q =
+            "select name, setting from pg_settings where name in  ('max_wal_senders', 'max_worker_processes', 'wal_level',  'max_replication_slots', 'max_logical_replication_workers');"
+
+          {
+            source_db_is_superuser: is_super_user?(source_db_url),
+            target_db_is_superuser: is_super_user?(target_db_url),
+            source_db:
+              Query.run(
+                query: q,
+                connection_url: source_db_url,
+                user: db_user(source_db_url),
+              ),
+            target_db:
+              Query.run(
+                query: q,
+                connection_url: target_db_url,
+                user: db_user(target_db_url),
+              ),
+          }
+        rescue => e
+          abort_with("Unable to check config: #{e.message}")
+        end
+    end
+
+    def assert_config
+      unless assert_wal_level_logical(config.dig(:source_db))
+        abort_with("WAL_LEVEL should be LOGICAL on source DB")
+      end
+
+      unless assert_wal_level_logical(config.dig(:target_db))
+        abort_with("WAL_LEVEL should be LOGICAL on target DB")
+      end
+
+      unless config.dig(:source_db_is_superuser)
+        abort_with("User on source database should be a superuser")
+      end
+
+      return if config.dig(:target_db_is_superuser)
+      abort_with("User on target database should be a superuser")
+    end
+
+    def bootstrap(options)
+      assert_config
+      logger.info("Setting up schema")
+      setup_schema
+
+      logger.info("Setting up replication user on source database")
+      create_user(conn_string: source_db_url, group_name: options[:group_name])
+
+      logger.info("Setting up replication user on target database")
+      create_user(conn_string: target_db_url, group_name: options[:group_name])
+
+      logger.info("Setting up groups tables")
+      Group.setup
+    rescue => e
+      abort_with("Unable to bootstrap: #{e.message}")
+    end
+
+    def cleanup(options)
+      logger.info("Dropping groups table")
+      Group.drop
+
+      if options[:everything]
+        logger.info("Dropping schema")
+        drop_schema
+      end
+
+      if options[:everything] || options[:sync]
+        Orchestrate.drop_publication(
+          group_name: options[:group_name],
+          conn_string: source_db_url,
+        )
+
+        Orchestrate.drop_subscription(
+          group_name: options[:group_name],
+          target_conn_string: target_db_url,
+        )
+      end
+
+      if options[:everything]
+        # Drop users at last
+        logger.info("Dropping replication user on source database")
+        drop_user(conn_string: source_db_url, group_name: options[:group_name])
+
+        logger.info("Dropping replication user on target database")
+        drop_user(conn_string: target_db_url, group_name: options[:group_name])
+      end
+    rescue => e
+      abort_with("Unable to cleanup: #{e.message}")
+    end
+
+    def drop_schema
+      Query.run(
+        query: "DROP SCHEMA IF EXISTS #{internal_schema_name} CASCADE",
+        connection_url: source_db_url,
+        schema: internal_schema_name,
+      )
+    end
+
+    def setup_schema
+      sql = <<~SQL
+        create schema if not exists #{internal_schema_name};
+        grant usage on schema #{internal_schema_name} to #{db_user(source_db_url)};
+        grant create on schema #{internal_schema_name} to #{db_user(source_db_url)};
+      SQL
+
+      Query.run(
+        query: sql,
+        connection_url: source_db_url,
+        schema: internal_schema_name,
+        user: db_user(target_db_url),
+      )
+    end
+
+    def logger
+      @logger ||=
+        begin
+          logger = Ougai::Logger.new($stdout)
+          logger.level =
+            ENV["DEBUG"] ? Ougai::Logger::TRACE : Ougai::Logger::INFO
+          logger.with_fields = { version: PgEasyReplicate::VERSION }
+          logger
+        end
+    end
+
+    private
+
+    def assert_wal_level_logical(db_config)
+      db_config&.find do |r|
+        r.dig(:name) == "wal_level" && r.dig(:setting) == "logical"
+      end
+    end
+
+    def is_super_user?(url)
+      Query.run(
+        query:
+          "select usesuper from pg_user where usename = '#{db_user(url)}';",
+        connection_url: url,
+        user: db_user(target_db_url),
+      ).first[
+        :usesuper
+      ]
+    end
+
+    def create_user(conn_string:, group_name:)
+      password = connection_info(conn_string)[:user]
+      sql = <<~SQL
+        drop role if exists #{internal_user_name};
+        create role #{internal_user_name} with password '#{password}' login superuser createdb createrole;
+      SQL
+
+      Query.run(
+        query: sql,
+        connection_url: conn_string,
+        user: db_user(target_db_url),
+      )
+    end
+
+    def drop_user(conn_string:, group_name:)
+      sql = "drop role if exists #{internal_user_name};"
+      Query.run(
+        query: sql,
+        connection_url: conn_string,
+        user: db_user(conn_string),
+      )
+    end
+  end
+end

data/package.json

@@ -0,0 +1,13 @@
+{
+  "name": "pg-osc",
+  "version": "1.0.0",
+  "main": "index.js",
+  "repository": "git@github.com:shayonj/pg-osc.git",
+  "author": "Shayon Mukherjee <shayonj@gmail.com>",
+  "license": "MIT",
+  "private": true,
+  "dependencies": {
+    "@prettier/plugin-ruby": "^3.2.2",
+    "prettier": "^2.8.8"
+  }
+}