From 35a449e60370cefbca7a5bf281e0a6b4599d6952 Mon Sep 17 00:00:00 2001
From: James Cor <jcor@ucsd.edu>
Date: Tue, 19 Nov 2024 14:10:55 -0800
Subject: [PATCH] support `EXISTS` subquery (#978)

---
 server/ast/expr.go                     |  2 +-
 server/ast/subquery.go                 | 19 ++++++++---
 testing/go/regression/tool/replay.go   |  2 +-
 testing/go/regression/tool/run_test.go |  1 +
 testing/go/subqueries_test.go          | 46 ++++++++++++++++++++++++++
 5 files changed, 63 insertions(+), 7 deletions(-)

diff --git a/server/ast/expr.go b/server/ast/expr.go
index 8ee33f18e1..1e1e7a2a1a 100644
--- a/server/ast/expr.go
+++ b/server/ast/expr.go
@@ -727,7 +727,7 @@ func nodeExpr(ctx *Context, node tree.Expr) (vitess.Expr, error) {
 			Expression: unknownLiteral,
 		}, nil
 	case *tree.Subquery:
-		return nodeSubquery(ctx, node)
+		return nodeSubqueryOrExists(ctx, node)
 	case *tree.Tuple:
 		if len(node.Labels) > 0 {
 			return nil, fmt.Errorf("tuple labels are not yet supported")
diff --git a/server/ast/subquery.go b/server/ast/subquery.go
index 7ca20e16c1..02dff93e67 100644
--- a/server/ast/subquery.go
+++ b/server/ast/subquery.go
@@ -15,8 +15,6 @@
 package ast
 
 import (
-	"fmt"
-
 	vitess "github.com/dolthub/vitess/go/vt/sqlparser"
 
 	"github.com/dolthub/doltgresql/postgres/parser/sem/tree"
@@ -28,9 +26,6 @@ func nodeSubquery(ctx *Context, node *tree.Subquery) (*vitess.Subquery, error) {
 	if node == nil {
 		return nil, nil
 	}
-	if node.Exists {
-		return nil, fmt.Errorf("EXISTS is not yet supported")
-	}
 	selectStmt, err := nodeSelectStatement(ctx, node.Select)
 	if err != nil {
 		return nil, err
@@ -51,3 +46,17 @@ func nodeSubqueryToTableExpr(ctx *Context, node *tree.Subquery) (vitess.TableExp
 		As:   vitess.NewTableIdent(utils.GenerateUniqueAlias()),
 	}, nil
 }
+
+// nodeSubqueryOrExists handles *tree.Subquery nodes that may be an EXISTS subquery, returning a vitess.Expr.
+func nodeSubqueryOrExists(ctx *Context, node *tree.Subquery) (vitess.Expr, error) {
+	subquery, err := nodeSubquery(ctx, node)
+	if err != nil {
+		return nil, err
+	}
+	if !node.Exists {
+		return subquery, nil
+	}
+	return &vitess.ExistsExpr{
+		Subquery: subquery,
+	}, nil
+}
diff --git a/testing/go/regression/tool/replay.go b/testing/go/regression/tool/replay.go
index 55537ac09c..53ec6f9a87 100644
--- a/testing/go/regression/tool/replay.go
+++ b/testing/go/regression/tool/replay.go
@@ -454,7 +454,7 @@ ListenerLoop:
 					}
 				}
 				for _, failQuery := range options.FailQueries {
-					if message.String == failQuery {
+					if strings.Contains(message.String, failQuery) {
 						tracker.Failed++
 						tracker.AddFailure(ReplayTrackerItem{
 							Query:           message.String,
diff --git a/testing/go/regression/tool/run_test.go b/testing/go/regression/tool/run_test.go
index 0f7eaade3b..696d3fd7a6 100644
--- a/testing/go/regression/tool/run_test.go
+++ b/testing/go/regression/tool/run_test.go
@@ -102,4 +102,5 @@ WHERE pg_class.oid=indexrelid
 	AND indrelid=pg_class_2.oid
 	AND pg_class_2.relname = 'clstr_tst'
 	AND indisclustered;`,
+	`SELECT 1 FROM pg_catalog.pg_constraint WHERE conrelid = i.indrelid AND conindid = i.indexrelid`,
 }
diff --git a/testing/go/subqueries_test.go b/testing/go/subqueries_test.go
index 7be234e1c9..beb7a33fcd 100755
--- a/testing/go/subqueries_test.go
+++ b/testing/go/subqueries_test.go
@@ -161,3 +161,49 @@ ORDER BY 1;`,
 		},
 	})
 }
+
+func TestExistSubquery(t *testing.T) {
+	RunScripts(t, []ScriptTest{
+		{
+			Name: "basic case",
+			SetUpScript: []string{
+				`CREATE TABLE test (id INT PRIMARY KEY);`,
+				`INSERT INTO test VALUES (1), (3), (2);`,
+			},
+			Assertions: []ScriptTestAssertion{
+				{
+					Query: `SELECT * FROM test WHERE EXISTS (SELECT 123);`,
+					Expected: []sql.Row{
+						{1},
+						{2},
+						{3},
+					},
+				},
+				{
+					Query:    `SELECT * FROM test WHERE NOT EXISTS (SELECT 123);`,
+					Expected: []sql.Row{},
+				},
+				{
+					Query: `SELECT 123 WHERE EXISTS (SELECT * FROM test);`,
+					Expected: []sql.Row{
+						{123},
+					},
+				},
+				{
+					Query:    `SELECT 123 WHERE EXISTS (SELECT * FROM test WHERE id > 10);`,
+					Expected: []sql.Row{},
+				},
+				{
+					Query:    `SELECT 123 WHERE NOT EXISTS (SELECT * FROM test);`,
+					Expected: []sql.Row{},
+				},
+				{
+					Query: `SELECT 123 WHERE NOT EXISTS (SELECT * FROM test WHERE id > 10);`,
+					Expected: []sql.Row{
+						{123},
+					},
+				},
+			},
+		},
+	})
+}