feat: change the schema tracker to track all UDFs

Signed-off-by: Andres Taylor <[email protected]>
vitessio · Apr 17, 2024 · e31db38 · e31db38
1 parent 178e6e8
commit e31db38
Show file tree

Hide file tree

Showing 10 changed files with 1,305 additions and 268 deletions.
diff --git a/go/vt/proto/query/query.pb.go b/go/vt/proto/query/query.pb.go
diff --git a/go/vt/proto/query/query_vtproto.pb.go b/go/vt/proto/query/query_vtproto.pb.go
diff --git a/go/vt/vtgate/schema/tracker.go b/go/vt/vtgate/schema/tracker.go
@@ -164,10 +164,13 @@ func (t *Tracker) loadUDFs(conn queryservice.QueryService, target *querypb.Targe
 	t.mu.Lock()
 	defer t.mu.Unlock()
 
-	err := conn.GetSchema(t.ctx, target, querypb.SchemaTableType_UDF_AGGREGATE, nil, func(schemaRes *querypb.GetSchemaResponse) error {
+	err := conn.GetSchema(t.ctx, target, querypb.SchemaTableType_UDFS, nil, func(schemaRes *querypb.GetSchemaResponse) error {
 		var udfs []string
-		for name := range schemaRes.TableDefinition {
-			udfs = append(udfs, name)
+		for _, udf := range schemaRes.Udfs {
+			if !udf.Aggregating {
+				continue
+			}
+			udfs = append(udfs, udf.Name)
 		}
 
 		t.udfs[target.Keyspace] = udfs