From 40bcf28ef59aa9c65a88bb6904d67ea8e9b42f4f Mon Sep 17 00:00:00 2001
From: Jeremy Woertink <jeremywoertink@gmail.com>
Date: Sat, 10 Feb 2024 15:21:52 -0800
Subject: [PATCH] Adding create_sequence and drop_sequence migration helpers.
 Fixes #910

---
 .../create_sequence_statement_spec.cr         | 14 ++++++++++
 .../migrator/drop_sequence_statement_spec.cr  |  8 ++++++
 spec/avram/migrator/migration_spec.cr         | 23 ++++++++++++++++
 .../migrator/create_sequence_statement.cr     | 26 +++++++++++++++++++
 src/avram/migrator/drop_sequence_statement.cr | 22 ++++++++++++++++
 src/avram/migrator/statement_helpers.cr       | 11 ++++++++
 6 files changed, 104 insertions(+)
 create mode 100644 spec/avram/migrator/create_sequence_statement_spec.cr
 create mode 100644 spec/avram/migrator/drop_sequence_statement_spec.cr
 create mode 100644 src/avram/migrator/create_sequence_statement.cr
 create mode 100644 src/avram/migrator/drop_sequence_statement.cr

diff --git a/spec/avram/migrator/create_sequence_statement_spec.cr b/spec/avram/migrator/create_sequence_statement_spec.cr
new file mode 100644
index 000000000..2e6e0738f
--- /dev/null
+++ b/spec/avram/migrator/create_sequence_statement_spec.cr
@@ -0,0 +1,14 @@
+require "../../spec_helper"
+
+describe Avram::Migrator::CreateSequenceStatement do
+  it "generates correct CREATE SEQUENCE sql" do
+    statement = Avram::Migrator::CreateSequenceStatement.new(:accounts_number).build
+    statement.should eq "CREATE SEQUENCE accounts_number_seq OWNED BY NONE;"
+
+    statement = Avram::Migrator::CreateSequenceStatement.new(:accounts_number, if_not_exists: true).build
+    statement.should eq "CREATE SEQUENCE IF NOT EXISTS accounts_number_seq OWNED BY NONE;"
+
+    statement = Avram::Migrator::CreateSequenceStatement.new(:accounts_number, owned_by: "accounts.number").build
+    statement.should eq "CREATE SEQUENCE accounts_number_seq OWNED BY accounts.number;"
+  end
+end
diff --git a/spec/avram/migrator/drop_sequence_statement_spec.cr b/spec/avram/migrator/drop_sequence_statement_spec.cr
new file mode 100644
index 000000000..330bae75e
--- /dev/null
+++ b/spec/avram/migrator/drop_sequence_statement_spec.cr
@@ -0,0 +1,8 @@
+require "../../spec_helper"
+
+describe Avram::Migrator::DropSequenceStatement do
+  it "generates correct DROP SEQUENCE sql" do
+    statement = Avram::Migrator::DropSequenceStatement.new(:accounts_number).build
+    statement.should eq "DROP SEQUENCE IF EXISTS accounts_number_seq;"
+  end
+end
diff --git a/spec/avram/migrator/migration_spec.cr b/spec/avram/migrator/migration_spec.cr
index 86f99f7ab..356301d8b 100644
--- a/spec/avram/migrator/migration_spec.cr
+++ b/spec/avram/migrator/migration_spec.cr
@@ -78,6 +78,16 @@ class MigrationCreateTableIfNotExists::V995 < Avram::Migrator::Migration::V1
   end
 end
 
+class MigrationCreateAndDropSequences::V994 < Avram::Migrator::Migration::V1
+  def migrate
+    create_sequence(:accounts, :number)
+  end
+
+  def rollback
+    drop_sequence(:accounts, :number)
+  end
+end
+
 describe Avram::Migrator::Migration::V1 do
   it "executes statements in a transaction" do
     expect_raises Avram::FailedMigration do
@@ -140,6 +150,19 @@ describe Avram::Migrator::Migration::V1 do
       sql.should contain "ALTER TABLE IF EXISTS fake_things"
     end
   end
+
+  describe "sequences" do
+    it "creates the sequences" do
+      migration = MigrationCreateAndDropSequences::V994.new
+      migration.migrate
+      sql = migration.prepared_statements.join("\n")
+      sql.should contain "CREATE SEQUENCE IF NOT EXISTS accounts_number_seq OWNED BY accounts.number"
+
+      migration.rollback
+      sql = migration.prepared_statements.join("\n")
+      sql.should contain "DROP SEQUENCE IF EXISTS accounts_number_seq"
+    end
+  end
 end
 
 private def get_column_names(table_name)
diff --git a/src/avram/migrator/create_sequence_statement.cr b/src/avram/migrator/create_sequence_statement.cr
new file mode 100644
index 000000000..a8a7e5ba3
--- /dev/null
+++ b/src/avram/migrator/create_sequence_statement.cr
@@ -0,0 +1,26 @@
+# Builds an SQL statement for creating a sequence using the given name appending "_seq".
+# Additional options may be provided for extra customization.
+#
+# ### Usage
+#
+# ```
+# CreateSequenceStatement.new(:accounts_number, if_not_exists: true, owned_by: "accounts.number").build
+# # => "CREATE SEQUENCE accounts_number_seq;"
+# ```
+class Avram::Migrator::CreateSequenceStatement
+  private getter? if_not_exists : Bool = false
+
+  def initialize(@name : String | Symbol, *, @if_not_exists : Bool = false, @owned_by : String = "NONE")
+  end
+
+  def build
+    name = "#{@name}_seq"
+
+    String.build do |io|
+      io << "CREATE SEQUENCE "
+      io << "IF NOT EXISTS " if if_not_exists?
+      io << name
+      io << " OWNED BY #{@owned_by};"
+    end
+  end
+end
diff --git a/src/avram/migrator/drop_sequence_statement.cr b/src/avram/migrator/drop_sequence_statement.cr
new file mode 100644
index 000000000..2c77c4fea
--- /dev/null
+++ b/src/avram/migrator/drop_sequence_statement.cr
@@ -0,0 +1,22 @@
+# Builds an SQL statement for dropping a sequence using the given name.
+#
+# ### Usage
+#
+# ```
+# DropSequenceStatement.new(:accounts_number, if_not_exists: true, owned_by: "accounts.number").build
+# # => "CREATE SEQUENCE accounts_number_seq;"
+# ```
+class Avram::Migrator::DropSequenceStatement
+  def initialize(@name : String | Symbol)
+  end
+
+  def build
+    name = "#{@name}_seq"
+
+    String.build do |io|
+      io << "DROP SEQUENCE IF EXISTS "
+      io << name
+      io << ';'
+    end
+  end
+end
diff --git a/src/avram/migrator/statement_helpers.cr b/src/avram/migrator/statement_helpers.cr
index 5342169c6..605520e4c 100644
--- a/src/avram/migrator/statement_helpers.cr
+++ b/src/avram/migrator/statement_helpers.cr
@@ -93,4 +93,15 @@ module Avram::Migrator::StatementHelpers
   def drop_trigger(table_name : TableName, name : String)
     prepared_statements << Avram::Migrator::DropTriggerStatement.new(table_name, name).build
   end
+
+  def create_sequence(table : TableName, column : Symbol)
+    sequence_name = "#{table}_#{column}"
+    owned_by = "#{table}.#{column}"
+    prepared_statements << Avram::Migrator::CreateSequenceStatement.new(sequence_name, if_not_exists: true, owned_by: owned_by).build
+  end
+
+  def drop_sequence(table : TableName, column : Symbol)
+    sequence_name = "#{table}_#{column}"
+    prepared_statements << Avram::Migrator::DropSequenceStatement.new(sequence_name).build
+  end
 end