Lock primitive GetFields implementation

go/test/endtoend/preparestmt/stmt_methods_test.go

@@ -379,3 +379,39 @@ func TestSelectDBA(t *testing.T) {
 	}
 	require.Equal(t, 2, rowCount)
 }
+
+func TestSelectLock(t *testing.T) {
+	defer cluster.PanicHandler(t)
+	dbo := Connect(t)
+	defer dbo.Close()
+
+	// Get Lock
+	prepare, err := dbo.Prepare("select get_lock(?, ?)")
+	require.NoError(t, err)
+
+	rows, err := prepare.Query("a", 100000)
+	require.NoError(t, err)
+
+	var resultBytes sql.RawBytes
+	require.True(t, rows.Next(), "no rows found")
+	err = rows.Scan(&resultBytes)
+	require.NoError(t, err)
+	assert.Equal(t, "1", string(resultBytes))
+
+	// for connection to be reused.
+	err = rows.Close()
+	require.NoError(t, err)
+
+	// Release Lock
+	prepare, err = dbo.Prepare("select release_lock(?)")
+	require.NoError(t, err)
+
+	rows, err = prepare.Query("a")
+	require.NoError(t, err)
+	defer rows.Close()
+
+	require.True(t, rows.Next(), "no rows found")
+	err = rows.Scan(&resultBytes)
+	require.NoError(t, err)
+	assert.Equal(t, "1", string(resultBytes))
+}

go/vt/vtgate/engine/lock.go

@@ -38,6 +38,8 @@ type Lock struct {
 	// Query specifies the query to be executed.
 	Query string
 
+	FieldQuery string
+
 	noInputs
 
 	noTxNeeded
@@ -60,6 +62,10 @@ func (l *Lock) GetTableName() string {
 
 // TryExecute is part of the Primitive interface
 func (l *Lock) TryExecute(vcursor VCursor, bindVars map[string]*querypb.BindVariable, _ bool) (*sqltypes.Result, error) {
+	return l.execLock(vcursor, l.Query, bindVars)
+}
+
+func (l *Lock) execLock(vcursor VCursor, query string, bindVars map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
 	rss, _, err := vcursor.ResolveDestinations(l.Keyspace.Name, nil, []key.Destination{l.TargetDestination})
 	if err != nil {
 		return nil, err
@@ -68,11 +74,11 @@ func (l *Lock) TryExecute(vcursor VCursor, bindVars map[string]*querypb.BindVari
 		return nil, vterrors.Errorf(vtrpc.Code_FAILED_PRECONDITION, "lock query can be routed to single shard only: %v", rss)
 	}
 
-	query := &querypb.BoundQuery{
-		Sql:           l.Query,
+	boundQuery := &querypb.BoundQuery{
+		Sql:           query,
 		BindVariables: bindVars,
 	}
-	return vcursor.ExecuteLock(rss[0], query)
+	return vcursor.ExecuteLock(rss[0], boundQuery)
 }
 
 // TryStreamExecute is part of the Primitive interface
@@ -86,12 +92,13 @@ func (l *Lock) TryStreamExecute(vcursor VCursor, bindVars map[string]*querypb.Bi
 
 // GetFields is part of the Primitive interface
 func (l *Lock) GetFields(vcursor VCursor, bindVars map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
-	return nil, vterrors.New(vtrpc.Code_UNIMPLEMENTED, "not implements in lock primitive")
+	return l.execLock(vcursor, l.FieldQuery, bindVars)
 }
 
 func (l *Lock) description() PrimitiveDescription {
 	other := map[string]interface{}{
-		"Query": l.Query,
+		"Query":      l.Query,
+		"FieldQuery": l.FieldQuery,
 	}
 	return PrimitiveDescription{
 		OperatorType:      "Lock",

go/vt/vtgate/planbuilder/select.go

@@ -343,10 +343,12 @@ func buildLockingPrimitive(sel *sqlparser.Select, vschema ContextVSchema) (engin
 	if err != nil {
 		return nil, err
 	}
+	buf := sqlparser.NewTrackedBuffer(sqlparser.FormatImpossibleQuery).WriteNode(sel)
 	return &engine.Lock{
 		Keyspace:          ks,
 		TargetDestination: key.DestinationKeyspaceID{0},
 		Query:             sqlparser.String(sel),
+		FieldQuery:        buf.String(),
 	}, nil
 }
 

go/vt/vtgate/planbuilder/testdata/lock_cases.txt

@@ -10,6 +10,7 @@
       "Sharded": false
     },
     "TargetDestination": "KeyspaceID(00)",
+    "FieldQuery": "select get_lock('xyz', 10) from dual where 1 != 1",
     "Query": "select get_lock('xyz', 10) from dual"
   }
 }
@@ -27,11 +28,30 @@ Gen4 plan same as above
       "Sharded": false
     },
     "TargetDestination": "KeyspaceID(00)",
+    "FieldQuery": "select is_free_lock('xyz') from dual where 1 != 1",
     "Query": "select is_free_lock('xyz') from dual"
   }
 }
 Gen4 plan same as above
 
+# get_lock from dual prepare query
+"select get_lock(?, ?)"
+{
+  "QueryType": "SELECT",
+  "Original": "select get_lock(?, ?)",
+  "Instructions": {
+    "OperatorType": "Lock",
+    "Keyspace": {
+      "Name": "main",
+      "Sharded": false
+    },
+    "TargetDestination": "KeyspaceID(00)",
+    "FieldQuery": "select get_lock(:v1, :v2) from dual where 1 != 1",
+    "Query": "select get_lock(:v1, :v2) from dual"
+  }
+}
+Gen4 plan same as above
+
 # lock tables read
 "lock tables t as x read local"
 {