add gms tests for

Author: elianddb

enginetest/memory_engine_test.go

@@ -200,23 +200,51 @@ func TestSingleQueryPrepared(t *testing.T) {
 
 // Convenience test for debugging a single query. Unskip and set to the desired query.
 func TestSingleScript(t *testing.T) {
-	t.Skip()
+	// t.Skip()
 	var scripts = []queries.ScriptTest{
 		{
-			Name:        "AS OF propagates to nested CALLs",
-			SetUpScript: []string{},
+			// https://github.com/dolthub/dolt/issues/9873
+			Name:        "FOR UPDATE OF syntax support tests",
+			SetUpScript: []string{
+				"CREATE TABLE task_instance (id INT PRIMARY KEY, task_id VARCHAR(255), dag_id VARCHAR(255), run_id VARCHAR(255), state VARCHAR(50), queued_by_job_id INT)",
+				"CREATE TABLE job (id INT PRIMARY KEY, state VARCHAR(50))",
+				"CREATE TABLE dag_run (dag_id VARCHAR(255), run_id VARCHAR(255), state VARCHAR(50))",
+				"CREATE TABLE t (id INT PRIMARY KEY, name VARCHAR(50))",
+				"INSERT INTO task_instance VALUES (1, 'task1', 'dag1', 'run1', 'running', 1)",
+				"INSERT INTO job VALUES (1, 'running')",
+				"INSERT INTO dag_run VALUES ('dag1', 'run1', 'running')",
+				"INSERT INTO t VALUES (1, 'test')",
+			},
 			Assertions: []queries.ScriptTestAssertion{
 				{
-					Query: "create procedure create_proc() create table t (i int primary key, j int);",
-					Expected: []sql.Row{
-						{types.NewOkResult(0)},
-					},
+					Query: `SELECT task_instance.id, task_instance.task_id, task_instance.dag_id, task_instance.run_id
+FROM task_instance INNER JOIN job ON job.id = task_instance.queued_by_job_id INNER JOIN dag_run ON dag_run.dag_id = task_instance.dag_id AND dag_run.run_id = task_instance.run_id
+ WHERE task_instance.state IN ('running', 'queued', 'scheduled') AND NOT (job.state <=> 'running') AND dag_run.state = 'running' FOR UPDATE OF task_instance SKIP LOCKED`,
+					Expected: []sql.Row{},
 				},
 				{
-					Query: "call create_proc()",
-					Expected: []sql.Row{
-						{types.NewOkResult(0)},
-					},
+					Query:    "SELECT * FROM t FOR UPDATE",
+					Expected: []sql.Row{{1, "test"}},
+				},
+				{
+					Query:    "SELECT * FROM t FOR UPDATE OF t",
+					Expected: []sql.Row{{1, "test"}},
+				},
+				{
+					Query:    "SELECT * FROM t FOR UPDATE OF t SKIP LOCKED",
+					Expected: []sql.Row{{1, "test"}},
+				},
+				{
+					Query:    "SELECT * FROM t FOR UPDATE OF t NOWAIT",
+					Expected: []sql.Row{{1, "test"}},
+				},
+				{
+					Query:    "SELECT * FROM task_instance t1, job t2 FOR UPDATE OF t1, t2",
+					Expected: []sql.Row{{1, "task1", "dag1", "run1", "running", 1, 1, "running"}},
+				},
+				{
+					Query:       "SELECT * FROM t FOR UPDATE OF nonexistent_table",
+					ExpectedErr: true,
 				},
 			},
 		},

go.mod

@@ -42,3 +42,5 @@ require (
 )
 
 go 1.25
+
+replace github.com/dolthub/vitess => /workspace/vitess

sql/planbuilder/select.go

@@ -18,6 +18,7 @@ import (
 	"fmt"
 
 	ast "github.com/dolthub/vitess/go/vt/sqlparser"
+	"github.com/sirupsen/logrus"
 
 	"github.com/dolthub/go-mysql-server/sql"
 	"github.com/dolthub/go-mysql-server/sql/expression"
@@ -59,6 +60,14 @@ func (b *Builder) buildSelect(inScope *scope, s *ast.Select) (outScope *scope) {
 	// 5) Build top-level scopes, replacing aggregation and aliases with
 	//    projections from (4).
 	// 6) Finish with final target projections.
+	
+	// Process the Lock field for FOR UPDATE clauses
+	if s.Lock != "" {
+		logrus.Info("SELECT statement has Lock field:", s.Lock)
+		// For now, we'll just log the lock information
+		// In a full implementation, this would create a lock node
+		// that would be applied to the final query result
+	}
 	fromScope := b.buildFrom(inScope, s.From)
 	if cn, ok := fromScope.node.(sql.CommentedNode); ok && len(s.Comments) > 0 {
 		fromScope.node = cn.WithComment(string(s.Comments[0]))