feat: Add source plugin version to table metadata.

Author: marianogappa

internal/servers/plugin/v3/plugin.go

@@ -39,6 +39,8 @@ func (s *Server) GetTables(ctx context.Context, req *pb.GetTables_Request) (*pb.
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "failed to get tables: %v", err)
 	}
+	pluginVersion := s.Plugin.Version()
+	tables.SetSourcePluginVersion(pluginVersion)
 	schemas := tables.ToArrowSchemas()
 	encoded := make([][]byte, len(schemas))
 	for i, sc := range schemas {
@@ -203,6 +205,9 @@ func (s *Server) Sync(req *pb.Sync_Request, stream pb.Plugin_SyncServer) error {
 		pbMsg := &pb.Sync_Response{}
 		switch m := msg.(type) {
 		case *message.SyncMigrateTable:
+			pluginVersion := s.Plugin.Version()
+			m.Table.SourcePluginVersion = pluginVersion
+			m.Table.Relations.SetSourcePluginVersion(pluginVersion)
 			tableSchema := m.Table.ToArrowSchema()
 			schemaBytes, err := pb.SchemaToBytes(tableSchema)
 			if err != nil {

internal/servers/plugin/v3/plugin_test.go

@@ -50,6 +50,32 @@ func TestGetVersion(t *testing.T) {
 	}
 }
 
+func TestGetTables(t *testing.T) {
+	ctx := context.Background()
+	pluginVersion := "v1.2.3"
+	s := Server{
+		Plugin: plugin.NewPlugin("test", pluginVersion, memdb.NewMemDBClient),
+	}
+
+	_, err := s.Init(ctx, &pb.Init_Request{})
+	require.NoError(t, err)
+
+	res, err := s.GetTables(ctx, &pb.GetTables_Request{Tables: []string{"*"}})
+	require.NoError(t, err)
+	require.NotNil(t, res)
+	require.Greater(t, len(res.Tables), 0, "expected at least one table")
+
+	// Verify that the plugin version is included in the schema metadata
+	for _, tableBytes := range res.Tables {
+		sc, err := pb.NewSchemaFromBytes(tableBytes)
+		require.NoError(t, err)
+
+		version, found := sc.Metadata().GetValue(schema.MetadataTablePluginVersion)
+		require.True(t, found, "expected plugin version to be in schema metadata")
+		require.Equal(t, pluginVersion, version, "expected plugin version to match")
+	}
+}
+
 type mockSyncServer struct {
 	grpc.ServerStream
 	messages []*pb.Sync_Response

schema/arrow.go

@@ -29,6 +29,7 @@ const (
 	MetadataTableIsPaid            = "cq:table_paid"
 	MetadataTablePermissionsNeeded = "cq:table_permissions_needed"
 	MetadataTableSensitiveColumns  = "cq:table_sensitive_columns"
+	MetadataTablePluginVersion     = "cq:table_plugin_version"
 )
 
 type Schemas []*arrow.Schema
@@ -112,6 +113,24 @@ func ReplaceFieldInRecord(src arrow.RecordBatch, fieldName string, field arrow.A
 	return record, nil
 }
 
+// AddPluginVersionToSchema adds the plugin version metadata to an Arrow schema.
+// It returns a new schema with the version added to the metadata.
+func AddPluginVersionToSchema(sc *arrow.Schema, version string) *arrow.Schema {
+	if version == "" {
+		return sc
+	}
+	existingMD := sc.Metadata()
+	md := make(map[string]string)
+	for i := 0; i < existingMD.Len(); i++ {
+		key := existingMD.Keys()[i]
+		value := existingMD.Values()[i]
+		md[key] = value
+	}
+	md[MetadataTablePluginVersion] = version
+	newMetadata := arrow.MetadataFrom(md)
+	return arrow.NewSchema(sc.Fields(), &newMetadata)
+}
+
 func AddInternalColumnsToRecord(record arrow.RecordBatch, cqClientIDValue string) (arrow.RecordBatch, error) {
 	schema := record.Schema()
 	nRows := int(record.NumRows())

schema/table.go

@@ -116,6 +116,10 @@ type Table struct {
 	// with whether the table makes use of a paid API or not.
 	IsPaid bool `json:"is_paid"`
 
+	// SourcePluginVersion is the version of the source plugin that created this table.
+	// This is set when tables are retrieved from a plugin and is included in the Arrow schema metadata.
+	SourcePluginVersion string `json:"source_plugin_version,omitempty"`
+
 	// IgnorePKComponentsMismatchValidation is a flag that indicates if the table should skip validating usage of both primary key components and primary keys
 	IgnorePKComponentsMismatchValidation bool `json:"ignore_pk_components_mismatch_validation"`
 }
@@ -232,6 +236,9 @@ func NewTableFromArrowSchema(sc *arrow.Schema) (*Table, error) {
 	if isPaid, found := tableMD.GetValue(MetadataTableIsPaid); found {
 		table.IsPaid = isPaid == MetadataTrue
 	}
+	if sourcePluginVersion, found := tableMD.GetValue(MetadataTablePluginVersion); found {
+		table.SourcePluginVersion = sourcePluginVersion
+	}
 	return table, nil
 }
 
@@ -375,6 +382,14 @@ func (tt Tables) FilterDfsFunc(include, exclude func(*Table) bool, skipDependent
 	return filteredTables
 }
 
+// SetSourcePluginVersion sets the SourcePluginVersion on all tables recursively, including relations.
+func (tt Tables) SetSourcePluginVersion(version string) {
+	for _, t := range tt {
+		t.SourcePluginVersion = version
+		t.Relations.SetSourcePluginVersion(version)
+	}
+}
+
 func (tt Tables) ToArrowSchemas() Schemas {
 	flattened := tt.FlattenTables()
 	schemas := make(Schemas, len(flattened))
@@ -606,6 +621,9 @@ func (t *Table) ToArrowSchema() *arrow.Schema {
 	md[MetadataTablePermissionsNeeded] = string(asJSON)
 	asJSON, _ = json.Marshal(t.SensitiveColumns)
 	md[MetadataTableSensitiveColumns] = string(asJSON)
+	if t.SourcePluginVersion != "" {
+		md[MetadataTablePluginVersion] = t.SourcePluginVersion
+	}
 
 	schemaMd := arrow.MetadataFrom(md)
 	for i, c := range t.Columns {

schema/table_test.go

@@ -717,7 +717,8 @@ func TestTablesToAndFromArrow(t *testing.T) {
 			Parent: &Table{
 				Name: "parent_table",
 			},
-			IsIncremental: true,
+			IsIncremental:       true,
+			SourcePluginVersion: "v1.0.0",
 			Columns: []Column{
 				{Name: "bool", Type: arrow.FixedWidthTypes.Boolean},
 				{Name: "int", Type: arrow.PrimitiveTypes.Int64},
@@ -747,6 +748,48 @@ func TestTablesToAndFromArrow(t *testing.T) {
 	}
 }
 
+func TestSetSourcePluginVersion(t *testing.T) {
+	version := "v2.1.0"
+	tables := Tables{
+		&Table{
+			Name: "table1",
+			Relations: Tables{
+				&Table{
+					Name: "table2",
+					Relations: Tables{
+						&Table{Name: "table3"},
+					},
+				},
+			},
+		},
+		&Table{Name: "table4"},
+	}
+
+	tables.SetSourcePluginVersion(version)
+
+	// Verify all tables have the version set
+	flattened := tables.FlattenTables()
+	for _, table := range flattened {
+		if table.SourcePluginVersion != version {
+			t.Errorf("expected SourcePluginVersion to be %q, got %q for table %s", version, table.SourcePluginVersion, table.Name)
+		}
+	}
+
+	// Verify relations also have the version set
+	for _, table := range tables {
+		if table.SourcePluginVersion != version {
+			t.Errorf("expected SourcePluginVersion to be %q, got %q for table %s", version, table.SourcePluginVersion, table.Name)
+		}
+		if len(table.Relations) > 0 {
+			for _, rel := range table.Relations {
+				if rel.SourcePluginVersion != version {
+					t.Errorf("expected SourcePluginVersion to be %q, got %q for relation %s", version, rel.SourcePluginVersion, rel.Name)
+				}
+			}
+		}
+	}
+}
+
 func TestValidateDuplicateTables(t *testing.T) {
 	tests := []struct {
 		name   string