Repository: supabase/cli
Pull request #1124 (2023-05-19 05:35:52)

fix: handle begin atomic function body state
## What kind of change does this PR introduce?

Bug fix https://github.com/supabase/cli/issues/610

## What is the new behavior?

- Handles new function body syntax: https://dba.stackexchange.com/a/300512
- Handles multiple statements in create rule: https://www.postgresql.org/docs/15/sql-createrule.html

## Additional context

Add any other context or screenshots.


coveralls: ## Pull Request Test Coverage Report for [Build 5090173870](https://coveralls.io/builds/60343846)

* **22** of **22**   **(100.0%)**  changed or added relevant lines in **1** file are covered.
* **7** unchanged lines in **1** file lost coverage.
* Overall coverage increased (+**0.04%**) to **63.077%**

---


|  Files with Coverage Reduction | New Missed Lines | % |
| :-----|--------------|--: |
| [internal/gen/keys/keys.go](https://coveralls.io/builds/60343846/source?filename=internal%2Fgen%2Fkeys%2Fkeys.go#L94) | 7 | 7.94% |
<!-- | **Total:** | **7** |  | -->

|  Totals | [![Coverage Status](https://coveralls.io/builds/60343846/badge)](https://coveralls.io/builds/60343846) |
| :-- | --: |
| Change from base [Build 5088830081](https://coveralls.io/builds/60340453): |  0.04% |
| Covered Lines: | 4534 |
| Relevant Lines: | 7188 |

---
##### 💛  - [Coveralls](https://coveralls.io)


Patch:
diff --git a/internal/gen/types/typescript/typescript.go b/internal/gen/types/typescript/typescript.go
index e1e2c1e7..aefd41fb 100644
--- a/internal/gen/types/typescript/typescript.go
+++ b/internal/gen/types/typescript/typescript.go
@@ -103,7 +103,7 @@ func Run(ctx context.Context, useLocal bool, useLinked bool, projectId string, d
 			},
 			[]string{
 				"node",
-				"dist/server/app.js",
+				"dist/server/server.js",
 				"gen",
 				"types",
 				"typescript",
diff --git a/internal/utils/parser/state.go b/internal/utils/parser/state.go
index e3a07ebb..38556f37 100644
--- a/internal/utils/parser/state.go
+++ b/internal/utils/parser/state.go
@@ -2,10 +2,19 @@ package parser
 
 import (
 	"bytes"
+	"strings"
 	"unicode"
 	"unicode/utf8"
 )
 
+const (
+	// Omit BEGIN to allow arbitrary whitespaces between BEGIN and ATOMIC keywords.
+	// This can fail if ATOMIC is used as column name because it is not a reserved
+	// keyword: https://www.postgresql.org/docs/current/sql-keywords-appendix.html
+	BEGIN_ATOMIC = "ATOMIC"
+	END_ATOMIC   = "END"
+)
+
 type State interface {
 	// Return nil to emit token
 	Next(r rune, data []byte) State
@@ -32,6 +41,15 @@ func (s *ReadyState) Next(r rune, data []byte) State {
 	case ';':
 		// Emit token
 		return nil
+	case '(':
+		return &AtomicState{prev: s, delimiter: []byte{')'}}
+	case 'c':
+		fallthrough
+	case 'C':
+		offset := len(data) - len(BEGIN_ATOMIC)
+		if offset >= 0 && strings.ToUpper(string(data[offset:])) == BEGIN_ATOMIC {
+			return &AtomicState{prev: s, delimiter: []byte(END_ATOMIC)}
+		}
 	}
 	return s
 }
@@ -107,7 +125,7 @@ type DollarState struct {
 func (s *DollarState) Next(r rune, data []byte) State {
 	window := data[len(data)-len(s.delimiter):]
 	if bytes.Equal(window, s.delimiter) {
-		// Break out of block state
+		// Break out of dollar state
 		return &ReadyState{}
 	}
 	return s
@@ -143,3 +161,24 @@ type EscapeState struct{}
 func (s *EscapeState) Next(r rune, data []byte) State {
 	return &ReadyState{}
 }
+
+// Opened BEGIN ATOMIC function body
+type AtomicState struct {
+	prev      State
+	delimiter []byte
+}
+
+func (s *AtomicState) Next(r rune, data []byte) State {
+	// If we are in a quoted state, the current delimiter doesn't count.
+	if curr := s.prev.Next(r, data); curr != nil {
+		s.prev = curr
+	}
+	if _, ok := s.prev.(*ReadyState); ok {
+		window := data[len(data)-len(s.delimiter):]
+		// Treat delimiter as case insensitive
+		if strings.ToUpper(string(window)) == string(s.delimiter) {
+			return &ReadyState{}
+		}
+	}
+	return s
+}


Test patch:
diff --git a/internal/utils/parser/state_test.go b/internal/utils/parser/state_test.go
index cba7dbe3..63e07288 100644
--- a/internal/utils/parser/state_test.go
+++ b/internal/utils/parser/state_test.go
@@ -152,3 +152,60 @@ func TestDoubleQuote(t *testing.T) {
 		assert.ElementsMatch(t, []string{`CREATE POLICY "pet""name" on pets;`, " END"}, stats)
 	})
 }
+
+func TestParentheses(t *testing.T) {
+	t.Run("preserves parentheses", func(t *testing.T) {
+		sql := []string{
+			`CREATE RULE notify_me AS ON UPDATE TO mytable DO (NOTIFY mytable; SELECT "1");`,
+			`SELECT 1`,
+		}
+		stats, err := Split(strings.NewReader(strings.Join(sql, "")))
+		require.NoError(t, err)
+		assert.ElementsMatch(t, sql, stats)
+	})
+
+	t.Run("ignores literals", func(t *testing.T) {
+		sql := `CREATE RULE notify_me AS ON UPDATE TO mytable DO (-- )
+SELECT ')';
+-- )
+)
+`
+		stats, err := Split(strings.NewReader(sql))
+		require.NoError(t, err)
+		assert.ElementsMatch(t, []string{sql}, stats)
+	})
+}
+
+func TestBeginAtomic(t *testing.T) {
+	t.Run("inline body", func(t *testing.T) {
+		sql := `CREATE FUNCTION add(a integer, b integer) RETURNS integer
+LANGUAGE SQL
+IMMUTABLE STRICT PARALLEL SAFE
+RETURNS NULL ON NULL INPUT
+BEGIN ATOMIC
+	SELECT 'add';
+	SELECT a + b;
+END;`
+		stats, err := Split(strings.NewReader(sql))
+		require.NoError(t, err)
+		assert.ElementsMatch(t, []string{sql}, stats)
+	})
+
+	t.Run("case insenstive", func(t *testing.T) {
+		sql := "begin atomic; select 'end'; end"
+		stats, err := Split(strings.NewReader(sql))
+		require.NoError(t, err)
+		assert.ElementsMatch(t, []string{sql}, stats)
+	})
+
+	t.Run("ignores literals", func(t *testing.T) {
+		sql := `CREATE FUNCTION test() BEGIN
+ATOMIC-- END;
+-- END;
+END
+;`
+		stats, err := Split(strings.NewReader(sql))
+		require.NoError(t, err)
+		assert.ElementsMatch(t, []string{sql}, stats)
+	})
+}
diff --git a/internal/utils/parser/token_test.go b/internal/utils/parser/token_test.go
index f0be87db..566b1186 100644
--- a/internal/utils/parser/token_test.go
+++ b/internal/utils/parser/token_test.go
@@ -32,13 +32,14 @@ func TestSplit(t *testing.T) {
 		return nil
 	}))
 	require.Len(t, fixture, 18)
+	fixture = append(fixture, "\n")
 
 	sql, err := os.Open(filepath.Join(testdata, "all.sql"))
 	require.NoError(t, err)
 	stats, err := Split(sql)
 	require.NoError(t, err)
 
-	assert.ElementsMatch(t, fixture, stats[:len(fixture)])
+	assert.ElementsMatch(t, fixture, stats)
 }
 
 func TestSplitAndTrim(t *testing.T) {