feat: add UDF fields to the KeyspaceSchema

Signed-off-by: Andres Taylor <[email protected]>

go/test/vschemawrapper/vschema_wrapper.go

@@ -147,6 +147,10 @@ func (vw *VSchemaWrapper) KeyspaceError(keyspace string) error {
 	return nil
 }
 
+func (vw *VSchemaWrapper) GetAggregateUDFs() (udfs []string) {
+	return vw.V.GetAggregateUDFs()
+}
+
 func (vw *VSchemaWrapper) GetForeignKeyChecksState() *bool {
 	return vw.ForeignKeyChecksState
 }

go/vt/schemadiff/semantics.go

@@ -71,6 +71,10 @@ func (si *declarativeSchemaInformation) KeyspaceError(keyspace string) error {
 	return nil
 }
 
+func (si *declarativeSchemaInformation) GetAggregateUDFs() []string {
+	return nil
+}
+
 func (si *declarativeSchemaInformation) GetForeignKeyChecksState() *bool {
 	return nil
 }

go/vt/vtgate/planbuilder/plan_test.go

@@ -590,19 +590,16 @@ func loadSchema(t testing.TB, filename string, setCollation bool) *vindexes.VSch
 		t.Fatal(err)
 	}
 	vschema := vindexes.BuildVSchema(formal, sqlparser.NewTestParser())
-	if err != nil {
-		t.Fatal(err)
-	}
+	require.NoError(t, err)
 	for _, ks := range vschema.Keyspaces {
-		if ks.Error != nil {
-			t.Fatal(ks.Error)
-		}
+		require.NoError(t, ks.Error)
 
 		// adding view in user keyspace
 		if ks.Keyspace.Name == "user" {
-			if err = vschema.AddView(ks.Keyspace.Name, "user_details_view", "select user.id, user_extra.col from user join user_extra on user.id = user_extra.user_id", sqlparser.NewTestParser()); err != nil {
-				t.Fatal(err)
-			}
+			err = vschema.AddView(ks.Keyspace.Name, "user_details_view", "select user.id, user_extra.col from user join user_extra on user.id = user_extra.user_id", sqlparser.NewTestParser())
+			require.NoError(t, err)
+			err = vschema.AddUDF(ks.Keyspace.Name, "udf_aggr")
+			require.NoError(t, err)
 		}
 
 		// setting a default value to all the text columns in the tables of this keyspace

go/vt/vtgate/planbuilder/plancontext/vschema.go

@@ -93,6 +93,9 @@ type VSchema interface {
 
 	// StorePrepareData stores the prepared data in the session.
 	StorePrepareData(name string, v *vtgatepb.PrepareData)
+
+	// GetAggregateUDFs returns the list of aggregate UDFs.
+	GetAggregateUDFs() []string
 }
 
 // PlannerNameToVersion returns the numerical representation of the planner

go/vt/vtgate/semantics/FakeSI.go

@@ -36,6 +36,7 @@ type FakeSI struct {
 	VindexTables     map[string]vindexes.Vindex
 	KsForeignKeyMode map[string]vschemapb.Keyspace_ForeignKeyMode
 	KsError          map[string]error
+	UDFs             []string
 }
 
 // FindTableOrVindex implements the SchemaInformation interface
@@ -80,3 +81,7 @@ func (s *FakeSI) KeyspaceError(keyspace string) error {
 	}
 	return nil
 }
+
+func (s *FakeSI) GetAggregateUDFs() []string {
+	return s.UDFs
+}

go/vt/vtgate/semantics/early_rewriter_test.go

@@ -645,6 +645,7 @@ func getSchemaWithKnownColumns() *FakeSI {
 				ColumnListAuthoritative: true,
 			},
 		},
+		UDFs: []string{"custom_udf"},
 	}
 	return schemaInfo
 }

go/vt/vtgate/semantics/info_schema.go

@@ -1661,3 +1661,7 @@ func (i *infoSchemaWithColumns) GetForeignKeyChecksState() *bool {
 func (i *infoSchemaWithColumns) KeyspaceError(keyspace string) error {
 	return i.inner.KeyspaceError(keyspace)
 }
+
+func (i *infoSchemaWithColumns) GetAggregateUDFs() []string {
+	return i.inner.GetAggregateUDFs()
+}

go/vt/vtgate/semantics/semantic_state.go

@@ -161,6 +161,7 @@ type (
 		ForeignKeyMode(keyspace string) (vschemapb.Keyspace_ForeignKeyMode, error)
 		GetForeignKeyChecksState() *bool
 		KeyspaceError(keyspace string) error
+		GetAggregateUDFs() []string
 	}
 
 	shortCut = int

go/vt/vtgate/vcursor_impl.go

@@ -1075,6 +1075,10 @@ func (vc *vcursorImpl) KeyspaceError(keyspace string) error {
 	return ks.Error
 }
 
+func (vc *vcursorImpl) GetAggregateUDFs() []string {
+	return vc.vschema.GetAggregateUDFs()
+}
+
 // ParseDestinationTarget parses destination target string and sets default keyspace if possible.
 func parseDestinationTarget(targetString string, vschema *vindexes.VSchema) (string, topodatapb.TabletType, key.Destination, error) {
 	destKeyspace, destTabletType, dest, err := topoprotopb.ParseDestination(targetString, defaultTabletType)

go/vt/vtgate/vindexes/vschema.go

@@ -245,6 +245,9 @@ type KeyspaceSchema struct {
 	Views           map[string]sqlparser.SelectStatement
 	Error           error
 	MultiTenantSpec *vschemapb.MultiTenantSpec
+
+	// These are the UDFs that exist in the schema and are aggregations
+	AggregateUDFs []string
 }
 
 type ksJSON struct {
@@ -422,6 +425,18 @@ func (vschema *VSchema) AddView(ksname, viewName, query string, parser *sqlparse
 	return nil
 }
 
+// AddTable adds a table to an existing keyspace in the VSchema.
+// It's only used from tests.
+func (vschema *VSchema) AddUDF(ksname, udfName string) error {
+	ks, ok := vschema.Keyspaces[ksname]
+	if !ok {
+		return fmt.Errorf("keyspace %s not found in vschema", ksname)
+	}
+
+	ks.AggregateUDFs = append(ks.AggregateUDFs, udfName)
+	return nil
+}
+
 func buildGlobalTables(source *vschemapb.SrvVSchema, vschema *VSchema) {
 	for ksname, ks := range source.Keyspaces {
 		ksvschema := vschema.Keyspaces[ksname]
@@ -1272,6 +1287,20 @@ func (vschema *VSchema) ResetCreated() {
 	vschema.created = time.Time{}
 }
 
+func (vschema *VSchema) GetAggregateUDFs() (udfs []string) {
+	seen := make(map[string]bool)
+	for _, ks := range vschema.Keyspaces {
+		for _, udf := range ks.AggregateUDFs {
+			if seen[udf] {
+				continue
+			}
+			seen[udf] = true
+			udfs = append(udfs, udf)
+		}
+	}
+	return
+}
+
 // ByCost provides the interface needed for ColumnVindexes to
 // be sorted by cost order.
 type ByCost []*ColumnVindex