feat: support complex statements

Co-authored-by: Vlad Faust <[email protected]>

README.md

@@ -73,6 +73,18 @@ CREATE TABLE baz (
 -- Indexes
 CREATE UNIQUE INDEX baz_content_index ON baz (content);
 
+-- Statements which might contain semicolons
+-- +migrate start
+CREATE OR REPLACE FUNCTION trigger_set_timestamp()
+RETURNS TRIGGER AS
+$$
+BEGIN
+  NEW.updated_at = NOW();
+  RETURN NEW;
+END;
+$$ LANGUAGE plpgsql;
+-- +migrate end
+
 -- +migrate down
 DROP TABLE baz;
 ```

spec/migrate/migration_spec.cr

@@ -11,6 +11,18 @@ describe Migrate::Migration do
   -- Indexes
   CREATE UNIQUE INDEX foo_content_index ON foo (content);
 
+  -- Statements that might contain semicolons
+  -- +migrate start
+  CREATE OR REPLACE FUNCTION trigger_set_timestamp()
+  RETURNS TRIGGER AS
+  $$
+  BEGIN
+    NEW.updated_at = NOW();
+    RETURN NEW;
+  END;
+  $$ LANGUAGE plpgsql;
+  -- +migrate end
+
   -- +migrate down
   DROP TABLE foo;
 
@@ -30,8 +42,19 @@ describe Migrate::Migration do
     CREATE UNIQUE INDEX foo_content_index ON foo (content);
     SQL
 
+    q3 = <<-SQL
+    CREATE OR REPLACE FUNCTION trigger_set_timestamp()
+    RETURNS TRIGGER AS
+    $$
+    BEGIN
+      NEW.updated_at = NOW();
+      RETURN NEW;
+    END;
+    $$ LANGUAGE plpgsql;
+    SQL
+
     it do
-      migration.queries_up.should eq [q1, q2]
+      migration.queries_up.should eq [q1, q2, q3]
     end
   end
 

spec/migrations/10_create_baz.sql

@@ -7,5 +7,17 @@ CREATE TABLE baz (
 -- Indexes
 CREATE UNIQUE INDEX baz_content_index ON baz (content);
 
+-- Statements that might contain semicolons
+-- +migrate start
+CREATE OR REPLACE FUNCTION trigger_set_timestamp()
+RETURNS TRIGGER AS
+$$
+BEGIN
+  NEW.updated_at = NOW();
+  RETURN NEW;
+END;
+$$ LANGUAGE plpgsql;
+-- +migrate end
+
 -- +migrate down
 DROP TABLE baz;

src/migrate/migration.cr

@@ -10,6 +10,8 @@ module Migrate
 
     def initialize(lines : Iterator)
       direction : Direction? = nil
+      complex = false
+      complex_end = false
 
       buffer = String.new
 
@@ -22,6 +24,10 @@ module Migrate
             direction = Direction::Up
           when "down"
             direction = Direction::Down
+          when "start"
+            complex = true
+          when "end"
+            complex_end = true
           when "error"
             message = /^#{Regex.escape(CMD_PREFIX)} error (?<message>.+)$/.match(line).try &.["message"]
 
@@ -37,7 +43,10 @@ module Migrate
           buffer += (line + "\n") if direction
         end
 
-        if !line.starts_with?("--") && line.strip.ends_with?(";")
+        if (line.strip.ends_with?(";") && !complex) || complex_end
+          complex_end = false
+          complex = false
+
           case direction
           when Direction::Up
             @queries_up.push(buffer.dup.strip)