Ruby class to partition by day using native declarative partitioning in postgres

partition_by_day.rb

# https://www.postgresql.org/docs/13/ddl-partitioning.html#DDL-PARTITIONING-DECLARATIVE
class PartitionByDay
  class Row
    include Virtus.model

    attribute :name, String
    attribute :expression, String
  end

  def self.all(table)
    rows = SQL.hash_results <<-SQL, table: table
      SELECT pg_class.relname AS name,
        pg_get_expr(pg_class.relpartbound, pg_class.oid, true) AS expression
      FROM pg_class base_tb
        JOIN pg_inherits ON pg_inherits.inhparent = base_tb.oid
        JOIN pg_class ON pg_class.oid = pg_inherits.inhrelid
      WHERE base_tb.oid = :table::regclass;
    SQL

    rows.map { |row| Row.new(row) }
  end

  def self.create(attributes = {})
    table = attributes.fetch(:table)
    name = attributes.fetch(:name)
    from = attributes.fetch(:from)
    to = attributes.fetch(:to)

    return if exists?(name)

    binds = {
      table: SQL::LITERAL(table),
      name: SQL::LITERAL(name),
      from: from,
      to: to,
    }

    begin
      SQL.run <<~SQL, binds
        CREATE TABLE :name PARTITION OF :table FOR VALUES FROM (:from) TO (:to)
      SQL
    rescue => exception
      if exception.is_a?(PG::DuplicateTable) || exception.cause.is_a?(PG::DuplicateTable)
        # partition exists so move on with life
      else
        raise
      end
    end
  end

  def self.exists?(name)
    # TODO: cache this in memory for a period of time so we don't add extra
    # queries all the time (assuming AR doesn't cache)
    ActiveRecord::Base.connection.table_exists?(name)
  end
  class << self; alias exist? exists?; end

  def self.drop(name)
    return unless exists?(name)

    SQL.run <<~SQL, name: SQL::LITERAL(name)
      DROP TABLE :name
    SQL
  end

  def self.detach(table, name)
    return unless exists?(name)

    SQL.run <<~SQL, table: SQL::LITERAL(table), name: SQL::LITERAL(name)
      ALTER TABLE :table DETACH PARTITION :name;
    SQL
  end

  def initialize(table, time)
    raise ArgumentError, "time must not be nil" if time.nil?
    raise ArgumentError, "table must not be nil" if table.nil?

    @time = time
    @table = table
  end

  def name
    @name ||= "#{@table}_%d_%02d_%02d" % [@time.year, @time.month, @time.day]
  end

  def from
    @from ||= @time.utc.beginning_of_day
  end

  def to
    @to ||= @time.utc.tomorrow.beginning_of_day
  end

  def exists?
    self.class.exists?(name)
  end
  alias :exist? :exists?

  def create
    self.class.create(table: @table, name: name, from: from, to: to)
  end

  def drop
    self.class.drop(name)
  end

  def detach
    self.class.detach(@table, name)
  end
end

partition_by_day_test.rb

require 'test_helper'

class PartitionByDayTest < ActiveSupport::TestCase
  def table_name
    "partitioned_adapter_request_logs"
  end

  test "#name" do
    {
      Time.utc(2021, 9, 9, 12, 3, 3) => "#{table_name}_2021_09_09",
      Time.utc(2021, 10, 9, 12, 3, 3) => "#{table_name}_2021_10_09",
      Time.utc(2021, 10, 12, 12, 3, 3) => "#{table_name}_2021_10_12",
    }.each do |time, name|
      assert_equal name, PartitionByDay.new("partitioned_adapter_request_logs", time).name
      assert_equal name, PartitionByDay.new("partitioned_adapter_request_logs", time).name
    end
  end

  test "#from" do
    {
      Time.utc(2021, 9, 9, 12, 3, 3) => Time.utc(2021, 9, 9, 0, 0, 0),
      Time.utc(2021, 10, 12, 12, 3, 3) => Time.utc(2021, 10, 12, 0, 0, 0),
    }.each do |time, name|
      assert_equal name, PartitionByDay.new("partitioned_adapter_request_logs", time).from
      assert_equal name, PartitionByDay.new("partitioned_adapter_request_logs", time).from
    end
  end

  test "#to" do
    {
      Time.utc(2021, 9, 9, 12, 3, 3) => Time.utc(2021, 9, 10, 0, 0, 0),
      Time.utc(2021, 10, 12, 12, 3, 3) => Time.utc(2021, 10, 13, 0, 0, 0),
    }.each do |time, name|
      assert_equal name, PartitionByDay.new("partitioned_adapter_request_logs", time).to
      assert_equal name, PartitionByDay.new("partitioned_adapter_request_logs", time).to
    end
  end

  test "#create, #exists, #detach and #drop" do
    time = Time.utc(2021, 10, 12, 12, 3, 3)
    partition = PartitionByDay.new("partitioned_adapter_request_logs", time)
    refute_predicate partition, :exists?
    refute_predicate partition, :exist?
    refute PartitionByDay.exists?(partition.name)
    refute PartitionByDay.exist?(partition.name)

    partition.create
    assert_includes ActiveRecord::Base.connection.tables, partition.name
    assert_predicate partition, :exists?
    assert_predicate partition, :exist?
    assert PartitionByDay.exists?(partition.name)
    assert PartitionByDay.exist?(partition.name)

    partitions = PartitionByDay.all("partitioned_adapter_request_logs").map(&:name)
    assert_includes partitions, partition.name

    partition.detach
    partitions = PartitionByDay.all("partitioned_adapter_request_logs").map(&:name)
    refute_includes partitions, partition.name
    assert_predicate partition, :exists?
    assert_predicate partition, :exist?
    assert PartitionByDay.exists?(partition.name)
    assert PartitionByDay.exist?(partition.name)

    assert_includes ActiveRecord::Base.connection.tables, partition.name
    partition.drop
    refute_includes ActiveRecord::Base.connection.tables, partition.name
    refute_predicate partition, :exists?
    refute_predicate partition, :exist?
    refute PartitionByDay.exists?(partition.name)
    refute PartitionByDay.exist?(partition.name)
  end

  private

  def valid_attributes(attributes = {})
    {
      status: 200,
      method: "GET",
      path: "/adapter/features",
      ip: "127.0.0.1",
      token_id: tokens(:john_api).id,
      created_at: Time.zone.now,
      headers: {
        "content_type" => "application/json",
      },
    }.merge(attributes)
  end
end

partitioned_table_migration.rb

class CreatePartitionedAdapterRequestLogs < ActiveRecord::Migration[6.1]
  def change
    execute <<~SQL
      CREATE TABLE partitioned_adapter_request_logs (
        status smallint NOT NULL,
        method character varying(20) NOT NULL,
        path character varying(2000) NOT NULL,
        ip inet NOT NULL,
        token_id integer,
        created_at timestamp without time zone NOT NULL,
        headers jsonb
      ) PARTITION BY RANGE (created_at);
    SQL

    add_index :partitioned_adapter_request_logs, :token_id
    add_index :partitioned_adapter_request_logs, :created_at, order: {created_at: :desc}
  end
end