From e6a4ee728987d00c4246759607d5728fc2b1ff0d Mon Sep 17 00:00:00 2001 From: Andres Taylor Date: Mon, 15 Apr 2024 11:40:54 +0200 Subject: [PATCH] feat: add UDF fields to the KeyspaceSchema Signed-off-by: Andres Taylor --- go/test/vschemawrapper/vschema_wrapper.go | 4 +++ go/vt/schemadiff/semantics.go | 4 +++ go/vt/vtgate/planbuilder/plan_test.go | 15 ++++------ .../vtgate/planbuilder/plancontext/vschema.go | 3 ++ go/vt/vtgate/semantics/FakeSI.go | 5 ++++ go/vt/vtgate/semantics/early_rewriter_test.go | 1 + go/vt/vtgate/semantics/info_schema.go | 4 +++ go/vt/vtgate/semantics/semantic_state.go | 1 + go/vt/vtgate/vcursor_impl.go | 4 +++ go/vt/vtgate/vindexes/vschema.go | 29 +++++++++++++++++++ 10 files changed, 61 insertions(+), 9 deletions(-) diff --git a/go/test/vschemawrapper/vschema_wrapper.go b/go/test/vschemawrapper/vschema_wrapper.go index 0f8b47a9804..4d1c424dda8 100644 --- a/go/test/vschemawrapper/vschema_wrapper.go +++ b/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 } diff --git a/go/vt/schemadiff/semantics.go b/go/vt/schemadiff/semantics.go index ee175a37966..ccbf654f566 100644 --- a/go/vt/schemadiff/semantics.go +++ b/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 } diff --git a/go/vt/vtgate/planbuilder/plan_test.go b/go/vt/vtgate/planbuilder/plan_test.go index 6a2d0fe8b3f..8ffc88bb606 100644 --- a/go/vt/vtgate/planbuilder/plan_test.go +++ b/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 diff --git a/go/vt/vtgate/planbuilder/plancontext/vschema.go b/go/vt/vtgate/planbuilder/plancontext/vschema.go index 1bca80cbc94..8ac4c57bfd7 100644 --- a/go/vt/vtgate/planbuilder/plancontext/vschema.go +++ b/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 diff --git a/go/vt/vtgate/semantics/FakeSI.go b/go/vt/vtgate/semantics/FakeSI.go index 933f4cd40f8..1ca6718f1a8 100644 --- a/go/vt/vtgate/semantics/FakeSI.go +++ b/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 +} diff --git a/go/vt/vtgate/semantics/early_rewriter_test.go b/go/vt/vtgate/semantics/early_rewriter_test.go index 719d33b83d7..9577abd42aa 100644 --- a/go/vt/vtgate/semantics/early_rewriter_test.go +++ b/go/vt/vtgate/semantics/early_rewriter_test.go @@ -645,6 +645,7 @@ func getSchemaWithKnownColumns() *FakeSI { ColumnListAuthoritative: true, }, }, + UDFs: []string{"custom_udf"}, } return schemaInfo } diff --git a/go/vt/vtgate/semantics/info_schema.go b/go/vt/vtgate/semantics/info_schema.go index d7470e2fd0a..11e577f3fa7 100644 --- a/go/vt/vtgate/semantics/info_schema.go +++ b/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() +} diff --git a/go/vt/vtgate/semantics/semantic_state.go b/go/vt/vtgate/semantics/semantic_state.go index 1ea4bc2a889..6c6e495b33d 100644 --- a/go/vt/vtgate/semantics/semantic_state.go +++ b/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 diff --git a/go/vt/vtgate/vcursor_impl.go b/go/vt/vtgate/vcursor_impl.go index 15c6296f108..9372012f77d 100644 --- a/go/vt/vtgate/vcursor_impl.go +++ b/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) diff --git a/go/vt/vtgate/vindexes/vschema.go b/go/vt/vtgate/vindexes/vschema.go index 9e21505690c..49987aef4dc 100644 --- a/go/vt/vtgate/vindexes/vschema.go +++ b/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