From 0a926485c8d761043f7f09a06532ddfd74a2e4f2 Mon Sep 17 00:00:00 2001 From: erezrokah Date: Fri, 2 Oct 2026 12:23:19 +0300 Subject: [PATCH 1/2] feat: Add destination assessment interface and AssessTables handler --- examples/simple_plugin/go.mod | 2 +- examples/simple_plugin/go.sum | 4 +- go.mod | 2 +- go.sum | 4 +- internal/servers/plugin/v3/assess.go | 85 ++++++++++++ internal/servers/plugin/v3/assess_test.go | 156 ++++++++++++++++++++++ plugin/plugin_assess.go | 94 +++++++++++++ plugin/plugin_assess_test.go | 33 +++++ 8 files changed, 374 insertions(+), 6 deletions(-) create mode 100644 internal/servers/plugin/v3/assess.go create mode 100644 internal/servers/plugin/v3/assess_test.go create mode 100644 plugin/plugin_assess.go create mode 100644 plugin/plugin_assess_test.go diff --git a/examples/simple_plugin/go.mod b/examples/simple_plugin/go.mod index 8a63a0e3af..9328b09582 100644 --- a/examples/simple_plugin/go.mod +++ b/examples/simple_plugin/go.mod @@ -34,7 +34,7 @@ require ( github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/cloudquery/cloudquery-api-go v1.14.13 // indirect github.com/cloudquery/codegen v0.4.1 // indirect - github.com/cloudquery/plugin-pb-go v1.27.23 // indirect + github.com/cloudquery/plugin-pb-go v1.27.24-0.20261002091754-b8dcd1cbaee9 // indirect github.com/cloudquery/plugin-sdk/v2 v2.7.0 // indirect github.com/getsentry/sentry-go v0.49.0 // indirect github.com/ghodss/yaml v1.0.0 // indirect diff --git a/examples/simple_plugin/go.sum b/examples/simple_plugin/go.sum index f2da194fa9..78e2805ea1 100644 --- a/examples/simple_plugin/go.sum +++ b/examples/simple_plugin/go.sum @@ -58,8 +58,8 @@ github.com/cloudquery/cloudquery-api-go v1.14.13 h1:+lu1mLKqVSwrc02eqf9vfqGj9c0R github.com/cloudquery/cloudquery-api-go v1.14.13/go.mod h1:u4uzOBEss9hJZco/tULIZuObqiBD748kejlxZ2AqOxs= github.com/cloudquery/codegen v0.4.1 h1:c9D18N925tUvnDeGHIl3JWKj37TyII9daHufkf8hU+Y= github.com/cloudquery/codegen v0.4.1/go.mod h1:QWIOD6R1aCa+YM+th+9Qt9lZw+ztdJR9JDEMLWyazwM= -github.com/cloudquery/plugin-pb-go v1.27.23 h1:X08b+1rB1PKw4s+XFVlUUitTO0NrQ2oe9UHjJ/afSiA= -github.com/cloudquery/plugin-pb-go v1.27.23/go.mod h1:ZUJgTp6qCuurc/c8qnDe/U8eqQpWfDxTBsd/IK8DUeY= +github.com/cloudquery/plugin-pb-go v1.27.24-0.20261002091754-b8dcd1cbaee9 h1:mFYZDvEri1MMEWEqt5a9/FTXE8n/XXEOvs6HqzDOPo0= +github.com/cloudquery/plugin-pb-go v1.27.24-0.20261002091754-b8dcd1cbaee9/go.mod h1:06VOdOlk5Y64CahAxQnmGy6l7PXJYm3SAwNxpEWPwn4= github.com/cloudquery/plugin-sdk/v2 v2.7.0 h1:hRXsdEiaOxJtsn/wZMFQC9/jPfU1MeMK3KF+gPGqm7U= github.com/cloudquery/plugin-sdk/v2 v2.7.0/go.mod h1:pAX6ojIW99b/Vg4CkhnsGkRIzNaVEceYMR+Bdit73ug= github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= diff --git a/go.mod b/go.mod index 6ed947d930..1e139f0726 100644 --- a/go.mod +++ b/go.mod @@ -11,7 +11,7 @@ require ( github.com/bradleyjkemp/cupaloy/v2 v2.8.0 github.com/cloudquery/cloudquery-api-go v1.14.13 github.com/cloudquery/codegen v0.4.1 - github.com/cloudquery/plugin-pb-go v1.27.23 + github.com/cloudquery/plugin-pb-go v1.27.24-0.20261002091754-b8dcd1cbaee9 github.com/cloudquery/plugin-sdk/v2 v2.7.0 github.com/getsentry/sentry-go v0.49.0 github.com/goccy/go-json v0.10.6 diff --git a/go.sum b/go.sum index 7e665e5e13..b74b654e93 100644 --- a/go.sum +++ b/go.sum @@ -60,8 +60,8 @@ github.com/cloudquery/codegen v0.4.1 h1:c9D18N925tUvnDeGHIl3JWKj37TyII9daHufkf8h github.com/cloudquery/codegen v0.4.1/go.mod h1:QWIOD6R1aCa+YM+th+9Qt9lZw+ztdJR9JDEMLWyazwM= github.com/cloudquery/jsonschema v0.0.0-20260703174721-45e7e20e0ed8 h1:s7B+c57yTtVL8Zhcebae5poFInJwTftuMakj625CRbw= github.com/cloudquery/jsonschema v0.0.0-20260703174721-45e7e20e0ed8/go.mod h1:KMcD1TlufeD5r5DmbYvYWs+cyULzIVuKkd7jpoKWVjk= -github.com/cloudquery/plugin-pb-go v1.27.23 h1:X08b+1rB1PKw4s+XFVlUUitTO0NrQ2oe9UHjJ/afSiA= -github.com/cloudquery/plugin-pb-go v1.27.23/go.mod h1:ZUJgTp6qCuurc/c8qnDe/U8eqQpWfDxTBsd/IK8DUeY= +github.com/cloudquery/plugin-pb-go v1.27.24-0.20261002091754-b8dcd1cbaee9 h1:mFYZDvEri1MMEWEqt5a9/FTXE8n/XXEOvs6HqzDOPo0= +github.com/cloudquery/plugin-pb-go v1.27.24-0.20261002091754-b8dcd1cbaee9/go.mod h1:06VOdOlk5Y64CahAxQnmGy6l7PXJYm3SAwNxpEWPwn4= github.com/cloudquery/plugin-sdk/v2 v2.7.0 h1:hRXsdEiaOxJtsn/wZMFQC9/jPfU1MeMK3KF+gPGqm7U= github.com/cloudquery/plugin-sdk/v2 v2.7.0/go.mod h1:pAX6ojIW99b/Vg4CkhnsGkRIzNaVEceYMR+Bdit73ug= github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= diff --git a/internal/servers/plugin/v3/assess.go b/internal/servers/plugin/v3/assess.go new file mode 100644 index 0000000000..0bacd60a17 --- /dev/null +++ b/internal/servers/plugin/v3/assess.go @@ -0,0 +1,85 @@ +package plugin + +import ( + "context" + "fmt" + + pb "github.com/cloudquery/plugin-pb-go/pb/plugin/v3" + "github.com/cloudquery/plugin-sdk/v4/plugin" + "github.com/cloudquery/plugin-sdk/v4/schema" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +func (s *Server) AssessTables(ctx context.Context, req *pb.AssessTables_Request) (*pb.AssessTables_Response, error) { + tables := make([]plugin.TablePair, len(req.Tables)) + for i, pair := range req.Tables { + if len(pair.OldTable) == 0 && len(pair.NewTable) == 0 { + return nil, status.Errorf(codes.InvalidArgument, "table pair %d has neither an old nor a new table", i) + } + var err error + if tables[i].Old, err = tableFromBytes(pair.OldTable); err != nil { + return nil, status.Errorf(codes.InvalidArgument, "failed to decode old table: %v", err) + } + if tables[i].New, err = tableFromBytes(pair.NewTable); err != nil { + return nil, status.Errorf(codes.InvalidArgument, "failed to decode new table: %v", err) + } + } + findings, err := s.Plugin.AssessTables(ctx, tables, plugin.AssessOptions{MigrateForce: req.MigrateForce}) + if err != nil { + return nil, status.Errorf(codes.Internal, "failed to assess tables: %v", err) + } + resp := &pb.AssessTables_Response{Tables: make([]*pb.AssessTables_TableFinding, len(findings))} + for i, f := range findings { + resp.Tables[i] = tableFindingToPB(f) + } + return resp, nil +} + +func tableFromBytes(b []byte) (*schema.Table, error) { + if len(b) == 0 { + return nil, nil + } + sc, err := pb.NewSchemaFromBytes(b) + if err != nil { + return nil, err + } + table, err := schema.NewTableFromArrowSchema(sc) + if err != nil { + return nil, fmt.Errorf("failed to create table from schema: %w", err) + } + return table, nil +} + +func tableFindingToPB(f plugin.TableFinding) *pb.AssessTables_TableFinding { + columns := make([]*pb.AssessTables_ColumnFinding, len(f.Columns)) + for i, c := range f.Columns { + columns[i] = &pb.AssessTables_ColumnFinding{ + ColumnName: c.ColumnName, + Category: pb.AssessTables_Category(c.Category), + OldType: c.OldType, + NewType: c.NewType, + SafeModeBehavior: c.SafeModeBehavior, + ForcedModeBehavior: c.ForcedModeBehavior, + Evidence: evidenceToPB(c.Evidence), + } + } + return &pb.AssessTables_TableFinding{ + TableName: f.TableName, + Category: pb.AssessTables_Category(f.Category), + SafeModeBehavior: f.SafeModeBehavior, + ForcedModeBehavior: f.ForcedModeBehavior, + Columns: columns, + Evidence: evidenceToPB(f.Evidence), + CoverageIncomplete: f.CoverageIncomplete, + CoverageIncompleteReason: f.CoverageIncompleteReason, + } +} + +func evidenceToPB(evidence []plugin.Evidence) []*pb.AssessTables_Evidence { + res := make([]*pb.AssessTables_Evidence, len(evidence)) + for i, e := range evidence { + res[i] = &pb.AssessTables_Evidence{SyntheticValue: e.SyntheticValue, Before: e.Before, After: e.After} + } + return res +} diff --git a/internal/servers/plugin/v3/assess_test.go b/internal/servers/plugin/v3/assess_test.go new file mode 100644 index 0000000000..8ffed8b35d --- /dev/null +++ b/internal/servers/plugin/v3/assess_test.go @@ -0,0 +1,156 @@ +package plugin + +import ( + "context" + "net" + "testing" + + "github.com/apache/arrow-go/v18/arrow" + pb "github.com/cloudquery/plugin-pb-go/pb/plugin/v3" + "github.com/cloudquery/plugin-sdk/v4/plugin" + "github.com/cloudquery/plugin-sdk/v4/schema" + "github.com/google/go-cmp/cmp" + "github.com/rs/zerolog" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/status" + "google.golang.org/grpc/test/bufconn" + "google.golang.org/protobuf/testing/protocmp" +) + +type assessorClient struct { + mockSourceColumnAdderPluginClient + gotTables []plugin.TablePair + gotOptions plugin.AssessOptions +} + +func (c *assessorClient) AssessTables(_ context.Context, tables []plugin.TablePair, options plugin.AssessOptions) ([]plugin.TableFinding, error) { + c.gotTables, c.gotOptions = tables, options + return []plugin.TableFinding{{ + TableName: "test_table", + Category: plugin.AssessCategoryManualMigrationRequired, + SafeModeBehavior: "rejects this change", + ForcedModeBehavior: "drops and recreates the table", + Columns: []plugin.ColumnFinding{{ + ColumnName: "tags", + Category: plugin.AssessCategoryManualMigrationRequired, + OldType: "text[]", + NewType: "jsonb", + SafeModeBehavior: "rejects this change", + ForcedModeBehavior: "drops and recreates the table", + Evidence: []plugin.Evidence{{SyntheticValue: `["env:prod"]`, Before: `{"tags":["env:prod"]}`, After: `{"tags":["env:prod"]}`}}, + }}, + Evidence: []plugin.Evidence{{SyntheticValue: "header", Before: "tags", After: "tags"}}, + CoverageIncomplete: true, + CoverageIncompleteReason: "nested values not compared", + }}, nil +} + +func newAssessClient(t *testing.T, newClient plugin.NewClientFunc) pb.PluginClient { + t.Helper() + lis := bufconn.Listen(1024 * 1024) + srv := grpc.NewServer() + pb.RegisterPluginServer(srv, &Server{Plugin: plugin.NewPlugin("test", "development", newClient), Logger: zerolog.Nop()}) + go func() { _ = srv.Serve(lis) }() + t.Cleanup(srv.Stop) + + conn, err := grpc.NewClient("passthrough:///bufnet", + grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return lis.Dial() }), + grpc.WithTransportCredentials(insecure.NewCredentials())) + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + + client := pb.NewPluginClient(conn) + _, err = client.Init(context.Background(), &pb.Init_Request{NoConnection: true}) + require.NoError(t, err) + return client +} + +func tableBytes(t *testing.T, table *schema.Table) []byte { + t.Helper() + b, err := pb.SchemaToBytes(table.ToArrowSchema()) + require.NoError(t, err) + return b +} + +func TestAssessTablesRoundTrip(t *testing.T) { + assessor := &assessorClient{} + client := newAssessClient(t, func(context.Context, zerolog.Logger, []byte, plugin.NewClientOptions) (plugin.Client, error) { + return assessor, nil + }) + oldTable := &schema.Table{Name: "test_table", Columns: []schema.Column{{Name: "tags", Type: arrow.ListOf(arrow.BinaryTypes.String)}}} + newTable := &schema.Table{Name: "test_table", Columns: []schema.Column{{Name: "tags", Type: arrow.BinaryTypes.String}}} + + resp, err := client.AssessTables(context.Background(), &pb.AssessTables_Request{ + Tables: []*pb.AssessTables_TablePair{{OldTable: tableBytes(t, oldTable), NewTable: tableBytes(t, newTable)}, {NewTable: tableBytes(t, newTable)}}, + MigrateForce: true, + }) + require.NoError(t, err) + + require.True(t, assessor.gotOptions.MigrateForce) + require.Len(t, assessor.gotTables, 2) + require.Equal(t, "test_table", assessor.gotTables[0].Old.Name) + require.True(t, arrow.TypeEqual(arrow.ListOf(arrow.BinaryTypes.String), assessor.gotTables[0].Old.Columns[0].Type)) + require.True(t, arrow.TypeEqual(arrow.BinaryTypes.String, assessor.gotTables[0].New.Columns[0].Type)) + require.Nil(t, assessor.gotTables[1].Old) + + want := &pb.AssessTables_Response{Tables: []*pb.AssessTables_TableFinding{{ + TableName: "test_table", + Category: pb.AssessTables_CATEGORY_MANUAL_MIGRATION_REQUIRED, + SafeModeBehavior: "rejects this change", + ForcedModeBehavior: "drops and recreates the table", + Columns: []*pb.AssessTables_ColumnFinding{{ + ColumnName: "tags", + Category: pb.AssessTables_CATEGORY_MANUAL_MIGRATION_REQUIRED, + OldType: "text[]", + NewType: "jsonb", + SafeModeBehavior: "rejects this change", + ForcedModeBehavior: "drops and recreates the table", + Evidence: []*pb.AssessTables_Evidence{{SyntheticValue: `["env:prod"]`, Before: `{"tags":["env:prod"]}`, After: `{"tags":["env:prod"]}`}}, + }}, + Evidence: []*pb.AssessTables_Evidence{{SyntheticValue: "header", Before: "tags", After: "tags"}}, + CoverageIncomplete: true, + CoverageIncompleteReason: "nested values not compared", + }}} + require.Empty(t, cmp.Diff(want, resp, protocmp.Transform())) +} + +func TestAssessTablesWithoutAssessorReturnsUnknown(t *testing.T) { + client := newAssessClient(t, getColumnAdderPlugin()) + table := &schema.Table{Name: "test_table", Columns: []schema.Column{{Name: "id", Type: arrow.PrimitiveTypes.Int64}}} + + resp, err := client.AssessTables(context.Background(), &pb.AssessTables_Request{ + Tables: []*pb.AssessTables_TablePair{{OldTable: tableBytes(t, table)}}, + }) + require.NoError(t, err) + + want := &pb.AssessTables_Response{Tables: []*pb.AssessTables_TableFinding{{ + TableName: "test_table", + Category: pb.AssessTables_CATEGORY_UNKNOWN, + CoverageIncomplete: true, + CoverageIncompleteReason: plugin.AssessNotSupportedReason, + }}} + require.Empty(t, cmp.Diff(want, resp, protocmp.Transform())) +} + +func TestAssessTablesRejectsEmptyPair(t *testing.T) { + client := newAssessClient(t, getColumnAdderPlugin()) + _, err := client.AssessTables(context.Background(), &pb.AssessTables_Request{Tables: []*pb.AssessTables_TablePair{{}}}) + require.Equal(t, codes.InvalidArgument, status.Code(err)) +} + +func TestAssessCategoryMatchesProto(t *testing.T) { + for category, want := range map[plugin.AssessCategory]pb.AssessTables_Category{ + plugin.AssessCategoryUnknown: pb.AssessTables_CATEGORY_UNKNOWN, + plugin.AssessCategoryNoChange: pb.AssessTables_CATEGORY_NO_CHANGE, + plugin.AssessCategoryAutomaticallyMigratable: pb.AssessTables_CATEGORY_AUTOMATICALLY_MIGRATABLE, + plugin.AssessCategoryManualMigrationRequired: pb.AssessTables_CATEGORY_MANUAL_MIGRATION_REQUIRED, + plugin.AssessCategoryTableRemoved: pb.AssessTables_CATEGORY_TABLE_REMOVED, + plugin.AssessCategoryFileSchemaChanged: pb.AssessTables_CATEGORY_FILE_SCHEMA_CHANGED, + } { + require.Equal(t, want, tableFindingToPB(plugin.TableFinding{Category: category}).Category) + } + require.Len(t, pb.AssessTables_Category_name, 6) +} diff --git a/plugin/plugin_assess.go b/plugin/plugin_assess.go new file mode 100644 index 0000000000..65af7d881c --- /dev/null +++ b/plugin/plugin_assess.go @@ -0,0 +1,94 @@ +package plugin + +import ( + "context" + "errors" + + "github.com/cloudquery/plugin-sdk/v4/schema" +) + +type AssessCategory int + +const ( + AssessCategoryUnknown AssessCategory = iota + AssessCategoryNoChange + AssessCategoryAutomaticallyMigratable + AssessCategoryManualMigrationRequired + AssessCategoryTableRemoved + AssessCategoryFileSchemaChanged +) + +const AssessNotSupportedReason = "destination does not support assessment" + +// TablePair holds a table before and after a schema change. Old is nil for an added table, New is nil for a removed one. +type TablePair struct { + Old *schema.Table + New *schema.Table +} + +func (p TablePair) TableName() string { + if p.New != nil { + return p.New.Name + } + if p.Old != nil { + return p.Old.Name + } + return "" +} + +type AssessOptions struct { + MigrateForce bool +} + +type Evidence struct { + SyntheticValue string + Before string + After string +} + +type ColumnFinding struct { + ColumnName string + Category AssessCategory + OldType string + NewType string + SafeModeBehavior string + ForcedModeBehavior string + Evidence []Evidence +} + +type TableFinding struct { + TableName string + Category AssessCategory + SafeModeBehavior string + ForcedModeBehavior string + Columns []ColumnFinding + Evidence []Evidence + CoverageIncomplete bool + CoverageIncompleteReason string +} + +// Assessor is an optional DestinationClient interface that reports how schema changes would be applied, without writing anything. +// It is called after Init with NoConnection set, so implementations must not open database or cloud connections. +type Assessor interface { + AssessTables(ctx context.Context, tables []TablePair, options AssessOptions) ([]TableFinding, error) +} + +// AssessTables returns one Unknown finding per table when the client does not implement Assessor. +func (p *Plugin) AssessTables(ctx context.Context, tables []TablePair, options AssessOptions) ([]TableFinding, error) { + if p.client == nil { + return nil, errors.New("plugin not initialized. call Init() first") + } + if assessor, ok := p.client.(Assessor); ok { + return assessor.AssessTables(ctx, tables, options) + } + findings := make([]TableFinding, len(tables)) + for i, table := range tables { + findings[i] = TableFinding{ + TableName: table.TableName(), + Category: AssessCategoryUnknown, + CoverageIncomplete: true, + CoverageIncompleteReason: AssessNotSupportedReason, + } + } + return findings, nil +} diff --git a/plugin/plugin_assess_test.go b/plugin/plugin_assess_test.go new file mode 100644 index 0000000000..21007fa5e1 --- /dev/null +++ b/plugin/plugin_assess_test.go @@ -0,0 +1,33 @@ +package plugin + +import ( + "context" + "testing" + + "github.com/cloudquery/plugin-sdk/v4/schema" + "github.com/stretchr/testify/require" +) + +func TestAssessTablesWithoutAssessorReturnsUnknown(t *testing.T) { + ctx := context.Background() + p := NewPlugin("test", "v1.0.0", newTestPluginClient) + require.NoError(t, p.Init(ctx, nil, NewClientOptions{NoConnection: true})) + + findings, err := p.AssessTables(ctx, []TablePair{ + {Old: &schema.Table{Name: "changed"}, New: &schema.Table{Name: "changed"}}, + {New: &schema.Table{Name: "added"}}, + {Old: &schema.Table{Name: "removed"}}, + }, AssessOptions{}) + require.NoError(t, err) + + unknown := func(name string) TableFinding { + return TableFinding{TableName: name, Category: AssessCategoryUnknown, CoverageIncomplete: true, CoverageIncompleteReason: AssessNotSupportedReason} + } + require.Equal(t, []TableFinding{unknown("changed"), unknown("added"), unknown("removed")}, findings) +} + +func TestAssessTablesBeforeInit(t *testing.T) { + p := NewPlugin("test", "v1.0.0", newTestPluginClient) + _, err := p.AssessTables(context.Background(), nil, AssessOptions{}) + require.Error(t, err) +} From e33b75272f56f9b2bc56cb498b2c498bd27cb2ee Mon Sep 17 00:00:00 2001 From: erezrokah Date: Fri, 2 Oct 2026 16:43:02 +0300 Subject: [PATCH 2/2] feat: Add shared synthetic-record helpers for destination assessment --- schema/synthetic.go | 280 +++++++++++++++++++++++++++++++++++++++ schema/synthetic_test.go | 200 ++++++++++++++++++++++++++++ 2 files changed, 480 insertions(+) create mode 100644 schema/synthetic.go create mode 100644 schema/synthetic_test.go diff --git a/schema/synthetic.go b/schema/synthetic.go new file mode 100644 index 0000000000..a78ad58245 --- /dev/null +++ b/schema/synthetic.go @@ -0,0 +1,280 @@ +package schema + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "reflect" + "strings" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/cloudquery/plugin-sdk/v4/types" +) + +var ErrUnableToCompare = errors.New("unable to compare") + +type SyntheticCase string + +const ( + SyntheticCasePopulated SyntheticCase = "populated" + SyntheticCaseNull SyntheticCase = "null" + SyntheticCaseEmptyCollection SyntheticCase = "empty collection" + SyntheticCaseNested SyntheticCase = "nested" + SyntheticCaseEscaping SyntheticCase = "escaping" +) + +var syntheticCases = []SyntheticCase{ + SyntheticCasePopulated, + SyntheticCaseNull, + SyntheticCaseEmptyCollection, + SyntheticCaseNested, + SyntheticCaseEscaping, +} + +const syntheticEscapingString = `"comma, \"quote\", back\\slash\nnew line\ttab é"` + +// SyntheticPair is a value that means the same under the old and the new field. +// Value is its Arrow JSON form, which both fields accept and read back as equal values. +type SyntheticPair struct { + Case SyntheticCase + Value string +} + +// SyntheticPairs returns equivalent values for each case both fields support. +// It returns ErrUnableToCompare when no populated value means the same under both fields. +func SyntheticPairs(oldField, newField arrow.Field) ([]SyntheticPair, error) { + var pairs []SyntheticPair + for _, c := range syntheticCases { + if c == SyntheticCaseNull { + if oldField.Nullable && newField.Nullable { + pairs = append(pairs, SyntheticPair{Case: c, Value: "null"}) + } + continue + } + candidates := append(syntheticCandidates(oldField.Type, c), syntheticCandidates(newField.Type, c)...) + for _, candidate := range candidates { + if isEquivalent(oldField.Type, newField.Type, candidate) { + pairs = append(pairs, SyntheticPair{Case: c, Value: candidate}) + break + } + } + } + if len(pairs) == 0 || pairs[0].Case != SyntheticCasePopulated { + return nil, fmt.Errorf("%w: no equivalent value for %s and %s", ErrUnableToCompare, oldField.Type, newField.Type) + } + return pairs, nil +} + +// SyntheticRecords builds one-column records for the old and the new field, with one row per pair. +func SyntheticRecords(oldField, newField arrow.Field, pairs []SyntheticPair) (oldRecord, newRecord arrow.RecordBatch, err error) { + values := make([]string, len(pairs)) + for i, pair := range pairs { + values[i] = pair.Value + } + rows := "[" + strings.Join(values, ",") + "]" + oldRecord, err = syntheticRecord(oldField, rows) + if err != nil { + return nil, nil, err + } + newRecord, err = syntheticRecord(newField, rows) + if err != nil { + oldRecord.Release() + return nil, nil, err + } + return oldRecord, newRecord, nil +} + +func syntheticRecord(field arrow.Field, rows string) (arrow.RecordBatch, error) { + bldr := array.NewRecordBuilder(memory.DefaultAllocator, arrow.NewSchema([]arrow.Field{field}, nil)) + defer bldr.Release() + if err := bldr.Field(0).UnmarshalJSON([]byte(rows)); err != nil { + return nil, fmt.Errorf("failed to build synthetic values for field %s: %w", field.Name, err) + } + return bldr.NewRecordBatch(), nil +} + +func isEquivalent(oldType, newType arrow.DataType, value string) bool { + oldValue, ok := readBack(oldType, value) + if !ok { + return false + } + newValue, ok := readBack(newType, value) + return ok && reflect.DeepEqual(oldValue, newValue) +} + +func readBack(dataType arrow.DataType, value string) (any, bool) { + bldr := array.NewBuilder(memory.DefaultAllocator, dataType) + defer bldr.Release() + if err := bldr.UnmarshalJSON([]byte("[" + value + "]")); err != nil || bldr.Len() != 1 { + return nil, false + } + arr := bldr.NewArray() + defer arr.Release() + marshaled, err := json.Marshal(arr.GetOneForMarshal(0)) + if err != nil { + return nil, false + } + var decoded any + if err := json.NewDecoder(bytes.NewReader(marshaled)).Decode(&decoded); err != nil { + return nil, false + } + return decoded, true +} + +func syntheticCandidates(dataType arrow.DataType, c SyntheticCase) []string { + switch dt := dataType.(type) { + case *arrow.ListType: + return syntheticListCandidates(dt.Elem(), c) + case *arrow.LargeListType: + return syntheticListCandidates(dt.Elem(), c) + case *arrow.MapType: + if c == SyntheticCaseNested && !isCollection(dt.ItemType()) { + return nil + } + return syntheticListCandidates(dt.Elem(), c) + case *arrow.StructType: + return syntheticStructCandidates(dt, c) + case *types.JSONType: + return syntheticJSONCandidates(c) + } + if c == SyntheticCasePopulated { + return syntheticScalarCandidates(dataType) + } + if c == SyntheticCaseEscaping && isStringLike(dataType) { + return []string{syntheticEscapingString} + } + return nil +} + +func syntheticListCandidates(elem arrow.DataType, c SyntheticCase) []string { + switch c { + case SyntheticCaseEmptyCollection: + return []string{"[]"} + case SyntheticCaseNested: + if !isCollection(elem) { + return nil + } + return wrapInList(syntheticCandidates(elem, SyntheticCasePopulated)) + default: + return wrapInList(syntheticCandidates(elem, c)) + } +} + +func wrapInList(values []string) []string { + wrapped := make([]string, len(values)) + for i, v := range values { + wrapped[i] = "[" + v + "]" + } + return wrapped +} + +func syntheticStructCandidates(dt *arrow.StructType, c SyntheticCase) []string { + if c == SyntheticCaseEmptyCollection { + return nil + } + if c == SyntheticCaseNested && !hasCollectionField(dt) { + return nil + } + hasEscaping := false + fields := make([]string, dt.NumFields()) + for i, field := range dt.Fields() { + value := firstCandidate(field.Type, SyntheticCasePopulated) + if c == SyntheticCaseEscaping { + if escaping := firstCandidate(field.Type, SyntheticCaseEscaping); escaping != "" { + value = escaping + hasEscaping = true + } + } + if value == "" { + return nil + } + name, _ := json.Marshal(field.Name) + fields[i] = string(name) + ":" + value + } + if c == SyntheticCaseEscaping && !hasEscaping { + return nil + } + return []string{"{" + strings.Join(fields, ",") + "}"} +} + +func firstCandidate(dataType arrow.DataType, c SyntheticCase) string { + candidates := syntheticCandidates(dataType, c) + if len(candidates) == 0 { + return "" + } + return candidates[0] +} + +func hasCollectionField(dt *arrow.StructType) bool { + for _, field := range dt.Fields() { + if isCollection(field.Type) { + return true + } + } + return false +} + +func isCollection(dataType arrow.DataType) bool { + switch dataType.(type) { + case *arrow.ListType, *arrow.LargeListType, *arrow.MapType, *arrow.StructType, *types.JSONType: + return true + } + return false +} + +func syntheticJSONCandidates(c SyntheticCase) []string { + switch c { + case SyntheticCasePopulated: + return []string{`{"env":"prod"}`} + case SyntheticCaseEmptyCollection: + return []string{`{}`, `[]`} + case SyntheticCaseNested: + return []string{`{"env":{"tags":["prod"]}}`} + case SyntheticCaseEscaping: + return []string{`{"note":` + syntheticEscapingString + `}`} + } + return nil +} + +func isStringLike(dataType arrow.DataType) bool { + return arrow.TypeEqual(dataType, arrow.BinaryTypes.String) || arrow.TypeEqual(dataType, arrow.BinaryTypes.LargeString) +} + +func syntheticScalarCandidates(dataType arrow.DataType) []string { + switch { + case isStringLike(dataType): + return []string{`"env:prod"`} + case arrow.TypeEqual(dataType, types.ExtensionTypes.UUID): + return []string{`"6ba7b810-9dad-11d1-80b4-00c04fd430c8"`} + case arrow.TypeEqual(dataType, types.ExtensionTypes.Inet): + return []string{`"192.0.2.1/24"`} + case arrow.TypeEqual(dataType, types.ExtensionTypes.MAC): + return []string{`"00:00:5e:00:53:01"`} + case arrow.IsInteger(dataType.ID()), arrow.IsFloating(dataType.ID()), dataType.ID() == arrow.DURATION: + return []string{`42`} + case arrow.IsDecimal(dataType.ID()): + return []string{`"12.5"`} + } + switch dataType.ID() { + case arrow.BOOL: + return []string{`true`} + case arrow.BINARY, arrow.LARGE_BINARY: + return []string{`"ZW52OnByb2Q="`} + case arrow.TIMESTAMP: + return []string{`"2024-01-02T03:04:05Z"`} + case arrow.DATE32, arrow.DATE64: + return []string{`"2024-01-02"`} + case arrow.TIME32, arrow.TIME64: + return []string{`"03:04:05"`} + case arrow.INTERVAL_MONTHS: + return []string{`{"months":1}`} + case arrow.INTERVAL_DAY_TIME: + return []string{`{"days":1,"milliseconds":1}`} + case arrow.INTERVAL_MONTH_DAY_NANO: + return []string{`{"months":1,"days":1,"nanoseconds":1}`} + } + return nil +} diff --git a/schema/synthetic_test.go b/schema/synthetic_test.go new file mode 100644 index 0000000000..a293e0f8df --- /dev/null +++ b/schema/synthetic_test.go @@ -0,0 +1,200 @@ +package schema + +import ( + "encoding/json" + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/cloudquery/plugin-sdk/v4/types" + "github.com/stretchr/testify/require" +) + +func nullableField(dataType arrow.DataType) arrow.Field { + return arrow.Field{Name: "col", Type: dataType, Nullable: true} +} + +func TestSyntheticPairs(t *testing.T) { + stringList := arrow.ListOf(arrow.BinaryTypes.String) + tests := []struct { + name string + oldField arrow.Field + newField arrow.Field + want []SyntheticPair + }{ + { + name: "list of strings to json", + oldField: nullableField(stringList), + newField: nullableField(types.ExtensionTypes.JSON), + want: []SyntheticPair{ + {Case: SyntheticCasePopulated, Value: `["env:prod"]`}, + {Case: SyntheticCaseNull, Value: "null"}, + {Case: SyntheticCaseEmptyCollection, Value: `[]`}, + {Case: SyntheticCaseEscaping, Value: `[` + syntheticEscapingString + `]`}, + }, + }, + { + name: "json to list of strings", + oldField: nullableField(types.ExtensionTypes.JSON), + newField: nullableField(stringList), + want: []SyntheticPair{ + {Case: SyntheticCasePopulated, Value: `["env:prod"]`}, + {Case: SyntheticCaseNull, Value: "null"}, + {Case: SyntheticCaseEmptyCollection, Value: `[]`}, + {Case: SyntheticCaseEscaping, Value: `[` + syntheticEscapingString + `]`}, + }, + }, + { + name: "nested list to json", + oldField: nullableField(arrow.ListOf(stringList)), + newField: nullableField(types.ExtensionTypes.JSON), + want: []SyntheticPair{ + {Case: SyntheticCasePopulated, Value: `[["env:prod"]]`}, + {Case: SyntheticCaseNull, Value: "null"}, + {Case: SyntheticCaseEmptyCollection, Value: `[]`}, + {Case: SyntheticCaseNested, Value: `[["env:prod"]]`}, + {Case: SyntheticCaseEscaping, Value: `[[` + syntheticEscapingString + `]]`}, + }, + }, + { + name: "struct to json", + oldField: nullableField(arrow.StructOf(arrow.Field{Name: "tags", Type: stringList, Nullable: true})), + newField: nullableField(types.ExtensionTypes.JSON), + want: []SyntheticPair{ + {Case: SyntheticCasePopulated, Value: `{"tags":["env:prod"]}`}, + {Case: SyntheticCaseNull, Value: "null"}, + {Case: SyntheticCaseNested, Value: `{"tags":["env:prod"]}`}, + {Case: SyntheticCaseEscaping, Value: `{"tags":[` + syntheticEscapingString + `]}`}, + }, + }, + { + name: "string to uuid", + oldField: nullableField(arrow.BinaryTypes.String), + newField: nullableField(types.ExtensionTypes.UUID), + want: []SyntheticPair{ + {Case: SyntheticCasePopulated, Value: `"6ba7b810-9dad-11d1-80b4-00c04fd430c8"`}, + {Case: SyntheticCaseNull, Value: "null"}, + }, + }, + { + name: "int32 to int64", + oldField: nullableField(arrow.PrimitiveTypes.Int32), + newField: nullableField(arrow.PrimitiveTypes.Int64), + want: []SyntheticPair{ + {Case: SyntheticCasePopulated, Value: `42`}, + {Case: SyntheticCaseNull, Value: "null"}, + }, + }, + { + name: "timestamp precision change", + oldField: nullableField(arrow.FixedWidthTypes.Timestamp_s), + newField: nullableField(arrow.FixedWidthTypes.Timestamp_us), + want: []SyntheticPair{ + {Case: SyntheticCasePopulated, Value: `"2024-01-02T03:04:05Z"`}, + {Case: SyntheticCaseNull, Value: "null"}, + }, + }, + { + name: "map to map with escaping", + oldField: nullableField(arrow.MapOf(arrow.BinaryTypes.String, arrow.BinaryTypes.String)), + newField: nullableField(arrow.MapOf(arrow.BinaryTypes.String, arrow.BinaryTypes.LargeString)), + want: []SyntheticPair{ + {Case: SyntheticCasePopulated, Value: `[{"key":"env:prod","value":"env:prod"}]`}, + {Case: SyntheticCaseNull, Value: "null"}, + {Case: SyntheticCaseEmptyCollection, Value: `[]`}, + {Case: SyntheticCaseEscaping, Value: `[{"key":` + syntheticEscapingString + `,"value":` + syntheticEscapingString + `}]`}, + }, + }, + { + name: "map of lists is nested", + oldField: nullableField(arrow.MapOf(arrow.BinaryTypes.String, stringList)), + newField: nullableField(arrow.MapOf(arrow.BinaryTypes.String, stringList)), + want: []SyntheticPair{ + {Case: SyntheticCasePopulated, Value: `[{"key":"env:prod","value":["env:prod"]}]`}, + {Case: SyntheticCaseNull, Value: "null"}, + {Case: SyntheticCaseEmptyCollection, Value: `[]`}, + {Case: SyntheticCaseNested, Value: `[{"key":"env:prod","value":["env:prod"]}]`}, + {Case: SyntheticCaseEscaping, Value: `[{"key":` + syntheticEscapingString + `,"value":[` + syntheticEscapingString + `]}]`}, + }, + }, + { + name: "not null field has no null pair", + oldField: arrow.Field{Name: "col", Type: arrow.BinaryTypes.String}, + newField: nullableField(arrow.BinaryTypes.String), + want: []SyntheticPair{ + {Case: SyntheticCasePopulated, Value: `"env:prod"`}, + {Case: SyntheticCaseEscaping, Value: syntheticEscapingString}, + }, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got, err := SyntheticPairs(tc.oldField, tc.newField) + require.NoError(t, err) + require.Equal(t, tc.want, got) + }) + } +} + +func TestSyntheticPairs_UnableToCompare(t *testing.T) { + tests := []struct { + name string + oldField arrow.Field + newField arrow.Field + }{ + {name: "int to string", oldField: nullableField(arrow.PrimitiveTypes.Int64), newField: nullableField(arrow.BinaryTypes.String)}, + {name: "string to int", oldField: nullableField(arrow.BinaryTypes.String), newField: nullableField(arrow.PrimitiveTypes.Int64)}, + {name: "duration unit change", oldField: nullableField(arrow.FixedWidthTypes.Duration_s), newField: nullableField(arrow.FixedWidthTypes.Duration_ms)}, + {name: "string to list", oldField: nullableField(arrow.BinaryTypes.String), newField: nullableField(arrow.ListOf(arrow.BinaryTypes.String))}, + {name: "bool to int", oldField: nullableField(arrow.FixedWidthTypes.Boolean), newField: nullableField(arrow.PrimitiveTypes.Int64)}, + {name: "fixed size list", oldField: nullableField(arrow.FixedSizeListOf(2, arrow.BinaryTypes.String)), newField: nullableField(arrow.FixedSizeListOf(2, arrow.BinaryTypes.String))}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + _, err := SyntheticPairs(tc.oldField, tc.newField) + require.ErrorIs(t, err, ErrUnableToCompare) + }) + } +} + +func TestSyntheticRecords(t *testing.T) { + oldField := arrow.Field{Name: "tags", Type: arrow.ListOf(arrow.BinaryTypes.String), Nullable: true} + newField := arrow.Field{Name: "tags", Type: types.ExtensionTypes.JSON, Nullable: true} + pairs, err := SyntheticPairs(oldField, newField) + require.NoError(t, err) + + oldRecord, newRecord, err := SyntheticRecords(oldField, newField, pairs) + require.NoError(t, err) + defer oldRecord.Release() + defer newRecord.Release() + + require.True(t, oldRecord.Schema().Equal(arrow.NewSchema([]arrow.Field{oldField}, nil))) + require.True(t, newRecord.Schema().Equal(arrow.NewSchema([]arrow.Field{newField}, nil))) + require.EqualValues(t, len(pairs), oldRecord.NumRows()) + require.EqualValues(t, len(pairs), newRecord.NumRows()) + + for i, pair := range pairs { + oldValue, err := json.Marshal(oldRecord.Column(0).GetOneForMarshal(i)) + require.NoError(t, err) + newValue, err := json.Marshal(newRecord.Column(0).GetOneForMarshal(i)) + require.NoError(t, err) + require.JSONEq(t, pair.Value, string(oldValue), "old value for case %s", pair.Case) + require.JSONEq(t, pair.Value, string(newValue), "new value for case %s", pair.Case) + } + populated, err := json.Marshal(newRecord.Column(0).GetOneForMarshal(0)) + require.NoError(t, err) + require.Equal(t, `["env:prod"]`, string(populated)) +} + +func TestSyntheticPairs_SameTypeIsComparable(t *testing.T) { + for _, column := range TestTable("test", TestSourceOptions{}).Columns { + t.Run(column.Name, func(t *testing.T) { + field := column.ToArrowField() + pairs, err := SyntheticPairs(field, field) + require.NoError(t, err) + oldRecord, newRecord, err := SyntheticRecords(field, field, pairs) + require.NoError(t, err) + oldRecord.Release() + newRecord.Release() + }) + } +}