Repository: ariga/atlas (go)

Issue:
Cannot use `DO $$`, `BEGIN` or `EXCEPTION` keywords in migrations.
Hey,

I have been trying to wrap some statements in an `Exception` but am not having much luck.

This is what I am adding to my `schema.sql` file:

```
DO $$ BEGIN
    CREATE TYPE some-type AS ENUM ('some-val-1', 'some-val-2);
EXCEPTION
    WHEN duplicate_object THEN null;
END; $$
```

When I try to generate a migration based on this, I get the following error:

```
atlas migrate diff new_migration \
      --dir "file://configs/pgdb/migrations" \
      --to "file://configs/pgdb/schema.sql" \
      --dev-url "docker://postgres?search_path=public"
      
Error: sql/migrate: read migration directory state: sql/migrate: execute: executing statement "DO $$ BEGIN\n    CREATE TYPE some-type AS ENUM ('some-val-1', 'some-val-2');" from version "schema": pq: unterminated dollar-quoted string at or near "$$ BEGIN
    CREATE TYPE some-type AS ENUM ('some-val-1', 'some-val-2');"
make: *** [gen-migrate] Error 1
```
I also tried to run this without the `Do $$` block and I get the following error:

```
Error: sql/migrate: read migration directory state: sql/migrate: execute: executing statement "BEGIN\n    CREATE TYPE some-type AS ENUM ('some-val-1', 'some-val-2');" from version "schema": pq: syntax error at or near "CREATE"
make: *** [gen-migrate] Error 1
```
I guess this syntax is not supported by the SQL parser that you guys use?

I also don't see anything in the docs around support for these keywords, such as `DO $$`, `BEGIN`, `EXCEPTION`.

Can you confirm if this is the case? Is there anything else I can do in the meantime? Do you have a timeframe for when you might support this?


Pull request description:
sql/migrate: support dollar quoted string in scanner
Fixed #1530 

Patch:
diff --git a/sql/migrate/lex.go b/sql/migrate/lex.go
index 539b9421..7a453eb0 100644
--- a/sql/migrate/lex.go
+++ b/sql/migrate/lex.go
@@ -8,6 +8,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"regexp"
 	"strings"
 	"unicode"
 	"unicode/utf8"
@@ -90,6 +91,9 @@ func newLex(input string) (*lex, error) {
 	return l, nil
 }
 
+// Dollar-quoted string as defined by the PostgreSQL scanner.
+var reDollarQuote = regexp.MustCompile(`^\$([A-Za-zÈ-ÿ_][\wÈ-ÿ]*)*\$`)
+
 func (l *lex) stmt() (*Stmt, error) {
 	var (
 		depth int
@@ -132,6 +136,10 @@ Scan:
 			l.addPos(len(l.delim) - l.width)
 			text = l.input[:l.pos]
 			break Scan
+		case r == '$' && reDollarQuote.MatchString(l.input[l.pos-1:]):
+			if err := l.skipDollarQuote(); err != nil {
+				return nil, err
+			}
 		case r == '#':
 			l.comment("#", "\n")
 		case r == '-' && l.next() == '-':
@@ -178,6 +186,27 @@ func (l *lex) skipQuote(quote rune) error {
 	}
 }
 
+func (l *lex) skipDollarQuote() error {
+	m := reDollarQuote.FindString(l.input[l.pos-1:])
+	if m == "" {
+		return fmt.Errorf("unexpected dollar quote at position %d", l.pos)
+	}
+	l.addPos(len(m) - 1)
+	for {
+		switch r := l.next(); {
+		case r == eos:
+			// Fail only if a delimiter was not set.
+			if l.delim == "" {
+				return errors.New("unclosed dollar-quoted string")
+			}
+			return nil
+		case r == '$' && strings.HasPrefix(l.input[l.pos-1:], m):
+			l.addPos(len(m) - 1)
+			return nil
+		}
+	}
+}
+
 func (l *lex) comment(left, right string) {
 	i := strings.Index(l.input[l.pos:], right)
 	// Not a comment.


Test patch:
diff --git a/sql/migrate/testdata/lex/15_dollar_quote.sql b/sql/migrate/testdata/lex/15_dollar_quote.sql
new file mode 100644
index 00000000..201af985
--- /dev/null
+++ b/sql/migrate/testdata/lex/15_dollar_quote.sql
@@ -0,0 +1,127 @@
+create function tcl_date_week(int4,int4,int4) returns text as $$
+    return [clock format [clock scan "$2/$3/$1"] -format "%U"]
+$$ language pltcl immutable;
+create function tclsnitch() returns event_trigger language pltcl as $$
+  elog NOTICE "tclsnitch: $TG_event $TG_tag"
+$$;
+
+create function foobar() returns int language sql as $$select 1;$$;
+
+create function tcl_composite_arg_ref2(T_comp1) returns text as '
+    return $1(ref2)
+' language pltcl;
+
+create function tcl_error_handling_test(text) returns text
+language pltcl
+as $function$
+    if {[catch $1 err]} {
+        if {[lindex $::errorCode 0] != "POSTGRES"} {
+            return $err
+        }
+        array set myArray $::errorCode
+        unset myArray(POSTGRES)
+        unset -nocomplain myArray(funcname)
+        unset -nocomplain myArray(filename)
+        unset -nocomplain myArray(lineno)
+
+# A comment.
+set vals []
+    foreach {key} [lsort [array names myArray]] {
+set value [string map {"\n" "\n\t"} $myArray($key)]
+            lappend vals "$key: $value"
+        }
+        return [join $vals "\n"]
+    } else {
+        return "no error"
+    }
+$function$;
+
+create function tcl_spi_exec(
+    prepare boolean,
+    action text
+)
+    returns void language pltcl AS $function$
+set query "select * from (values (1,'foo'),(2,'bar'),(3,'baz')) v(col1,col2)"
+    if {$1 == "t"} {
+set prep [spi_prepare $query {}]
+    spi_execp -array A $prep {
+        elog NOTICE "col1 $A(col1), col2 $A(col2)"
+
+        switch $A(col1) {
+            2 {
+                elog NOTICE "action: $2"
+                switch $2 {
+                    break {
+                        break
+                    }
+                    continue {
+                        continue
+                    }
+                    return {
+                        return
+                    }
+                    error {
+                        error "error message"
+                    }
+                }
+                error "should not get here"
+            }
+        }
+    }
+} else {
+    spi_exec -array A $query {
+        elog NOTICE "col1 $A(col1), col2 $A(col2)"
+
+        switch $A(col1) {
+            2 {
+                elog NOTICE "action: $2"
+                switch $2 {
+                    break {
+                        break
+                    }
+                    continue {
+                        continue
+                    }
+                    return {
+                        return
+                    }
+                    error {
+                        error "error message"
+                    }
+                }
+                error "should not get here"
+            }
+        }
+    }
+}
+elog NOTICE "end of function"
+$function$;
+
+DO $$ BEGIN
+CREATE TYPE some-type AS ENUM ('some-val-1', 'some-val-2');
+EXCEPTION
+    WHEN duplicate_object THEN null;
+END; $$;
+
+DO $$DECLARE r record;
+BEGIN
+    FOR r IN SELECT table_schema, table_name FROM information_schema.tables
+             WHERE table_type = 'VIEW' AND table_schema = 'public'
+    LOOP
+        EXECUTE 'GRANT ALL ON ' || quote_ident(r.table_schema) || '.' || quote_ident(r.table_name) || ' TO webuser';
+    END LOOP;
+END$$;
+
+# Shorter examples.
+SELECT * FROM table WHERE name = $$John Doe$$;
+SELECT * FROM table WHERE name = $a1$John Doe$a1$;
+UPDATE table SET description = $$Lorem ipsum dolor sit amet$$ WHERE id = 123;
+INSERT INTO table (name, description) VALUES ($$Jane Smith$$, $$This is Jane's description$$);
+CREATE INDEX index ON table (c);
+CREATE TRIGGER trigger BEFORE INSERT ON table FOR EACH ROW EXECUTE FUNCTION function(params);
+SELECT COUNT(*) FROM table WHERE description ILIKE $$%foo%$$;
+CREATE FUNCTION my_function() RETURNS void AS $$
+BEGIN
+   -- function body here
+END;
+$$ LANGUAGE plpgsql;
\ No newline at end of file
diff --git a/sql/migrate/testdata/lex/15_dollar_quote.sql.golden b/sql/migrate/testdata/lex/15_dollar_quote.sql.golden
new file mode 100644
index 00000000..dbe356f7
--- /dev/null
+++ b/sql/migrate/testdata/lex/15_dollar_quote.sql.golden
@@ -0,0 +1,134 @@
+create function tcl_date_week(int4,int4,int4) returns text as $$
+    return [clock format [clock scan "$2/$3/$1"] -format "%U"]
+$$ language pltcl immutable;
+-- end --
+create function tclsnitch() returns event_trigger language pltcl as $$
+  elog NOTICE "tclsnitch: $TG_event $TG_tag"
+$$;
+-- end --
+create function foobar() returns int language sql as $$select 1;$$;
+-- end --
+create function tcl_composite_arg_ref2(T_comp1) returns text as '
+    return $1(ref2)
+' language pltcl;
+-- end --
+create function tcl_error_handling_test(text) returns text
+language pltcl
+as $function$
+    if {[catch $1 err]} {
+        if {[lindex $::errorCode 0] != "POSTGRES"} {
+            return $err
+        }
+        array set myArray $::errorCode
+        unset myArray(POSTGRES)
+        unset -nocomplain myArray(funcname)
+        unset -nocomplain myArray(filename)
+        unset -nocomplain myArray(lineno)
+
+# A comment.
+set vals []
+    foreach {key} [lsort [array names myArray]] {
+set value [string map {"\n" "\n\t"} $myArray($key)]
+            lappend vals "$key: $value"
+        }
+        return [join $vals "\n"]
+    } else {
+        return "no error"
+    }
+$function$;
+-- end --
+create function tcl_spi_exec(
+    prepare boolean,
+    action text
+)
+    returns void language pltcl AS $function$
+set query "select * from (values (1,'foo'),(2,'bar'),(3,'baz')) v(col1,col2)"
+    if {$1 == "t"} {
+set prep [spi_prepare $query {}]
+    spi_execp -array A $prep {
+        elog NOTICE "col1 $A(col1), col2 $A(col2)"
+
+        switch $A(col1) {
+            2 {
+                elog NOTICE "action: $2"
+                switch $2 {
+                    break {
+                        break
+                    }
+                    continue {
+                        continue
+                    }
+                    return {
+                        return
+                    }
+                    error {
+                        error "error message"
+                    }
+                }
+                error "should not get here"
+            }
+        }
+    }
+} else {
+    spi_exec -array A $query {
+        elog NOTICE "col1 $A(col1), col2 $A(col2)"
+
+        switch $A(col1) {
+            2 {
+                elog NOTICE "action: $2"
+                switch $2 {
+                    break {
+                        break
+                    }
+                    continue {
+                        continue
+                    }
+                    return {
+                        return
+                    }
+                    error {
+                        error "error message"
+                    }
+                }
+                error "should not get here"
+            }
+        }
+    }
+}
+elog NOTICE "end of function"
+$function$;
+-- end --
+DO $$ BEGIN
+CREATE TYPE some-type AS ENUM ('some-val-1', 'some-val-2');
+EXCEPTION
+    WHEN duplicate_object THEN null;
+END; $$;
+-- end --
+DO $$DECLARE r record;
+BEGIN
+    FOR r IN SELECT table_schema, table_name FROM information_schema.tables
+             WHERE table_type = 'VIEW' AND table_schema = 'public'
+    LOOP
+        EXECUTE 'GRANT ALL ON ' || quote_ident(r.table_schema) || '.' || quote_ident(r.table_name) || ' TO webuser';
+    END LOOP;
+END$$;
+-- end --
+SELECT * FROM table WHERE name = $$John Doe$$;
+-- end --
+SELECT * FROM table WHERE name = $a1$John Doe$a1$;
+-- end --
+UPDATE table SET description = $$Lorem ipsum dolor sit amet$$ WHERE id = 123;
+-- end --
+INSERT INTO table (name, description) VALUES ($$Jane Smith$$, $$This is Jane's description$$);
+-- end --
+CREATE INDEX index ON table (c);
+-- end --
+CREATE TRIGGER trigger BEFORE INSERT ON table FOR EACH ROW EXECUTE FUNCTION function(params);
+-- end --
+SELECT COUNT(*) FROM table WHERE description ILIKE $$%foo%$$;
+-- end --
+CREATE FUNCTION my_function() RETURNS void AS $$
+BEGIN
+   -- function body here
+END;
+$$ LANGUAGE plpgsql;
\ No newline at end of file