feat: Compare generated Parquet schemas between table versions

parquet/assess.go

@@ -0,0 +1,98 @@
+package parquet
+
+import (
+	"strconv"
+	"strings"
+
+	"github.com/apache/arrow-go/v18/parquet"
+	"github.com/apache/arrow-go/v18/parquet/pqarrow"
+	pqschema "github.com/apache/arrow-go/v18/parquet/schema"
+	"github.com/cloudquery/plugin-sdk/v4/plugin"
+	"github.com/cloudquery/plugin-sdk/v4/schema"
+)
+
+// ParquetSchema returns the Parquet schema WriteHeader writes for the table, without writing anything.
+func (c *Client) ParquetSchema(t *schema.Table) (*pqschema.Schema, error) {
+	return pqarrow.ToParquet(convertSchema(t.ToArrowSchema()), c.writerProps(), pqarrow.DefaultWriterProps())
+}
+
+// AssessTable compares the Parquet schemas generated for the old and new table.
+func (c *Client) AssessTable(pair plugin.TablePair) (plugin.TableFinding, error) {
+	finding := plugin.TableFinding{TableName: pair.TableName(), Category: plugin.AssessCategoryNoChange}
+	if pair.New == nil {
+		finding.Category = plugin.AssessCategoryTableRemoved
+		return finding, nil
+	}
+	oldNames, oldTypes, err := c.fieldTypes(pair.Old)
+	if err != nil {
+		return plugin.TableFinding{}, err
+	}
+	newNames, newTypes, err := c.fieldTypes(pair.New)
+	if err != nil {
+		return plugin.TableFinding{}, err
+	}
+	for _, name := range oldNames {
+		if oldTypes[name] != newTypes[name] {
+			finding.Columns = append(finding.Columns, fileSchemaChange(name, oldTypes[name], newTypes[name]))
+		}
+	}
+	for _, name := range newNames {
+		if _, ok := oldTypes[name]; !ok {
+			finding.Columns = append(finding.Columns, fileSchemaChange(name, "", newTypes[name]))
+		}
+	}
+	if len(finding.Columns) > 0 {
+		finding.Category = plugin.AssessCategoryFileSchemaChanged
+	}
+	return finding, nil
+}
+
+func fileSchemaChange(name, oldType, newType string) plugin.ColumnFinding {
+	return plugin.ColumnFinding{
+		ColumnName: name,
+		Category:   plugin.AssessCategoryFileSchemaChanged,
+		OldType:    oldType,
+		NewType:    newType,
+	}
+}
+
+func (c *Client) fieldTypes(t *schema.Table) ([]string, map[string]string, error) {
+	typesByName := make(map[string]string)
+	if t == nil {
+		return nil, typesByName, nil
+	}
+	sc, err := c.ParquetSchema(t)
+	if err != nil {
+		return nil, nil, err
+	}
+	root := sc.Root()
+	names := make([]string, root.NumFields())
+	for i := range names {
+		field := root.Field(i)
+		names[i] = field.Name()
+		typesByName[field.Name()] = parquetType(field)
+	}
+	return names, typesByName, nil
+}
+
+func parquetType(node pqschema.Node) string {
+	repetition := strings.ToLower(node.RepetitionType().String())
+	logicalType := ""
+	if lt := node.LogicalType(); lt != nil && !lt.IsNone() {
+		logicalType = " (" + lt.String() + ")"
+	}
+	group, ok := node.(*pqschema.GroupNode)
+	if !ok {
+		primitive := node.(*pqschema.PrimitiveNode)
+		physicalType := strings.ToLower(primitive.PhysicalType().String())
+		if primitive.PhysicalType() == parquet.Types.FixedLenByteArray {
+			physicalType += "(" + strconv.Itoa(primitive.TypeLength()) + ")"
+		}
+		return repetition + " " + physicalType + logicalType
+	}
+	children := make([]string, group.NumFields())
+	for i := range children {
+		children[i] = group.Field(i).Name() + ": " + parquetType(group.Field(i))
+	}
+	return repetition + " group" + logicalType + " {" + strings.Join(children, "; ") + "}"
+}

parquet/assess_test.go

@@ -0,0 +1,172 @@
+package parquet
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/apache/arrow-go/v18/arrow"
+	"github.com/apache/arrow-go/v18/parquet/file"
+	"github.com/cloudquery/plugin-sdk/v4/plugin"
+	"github.com/cloudquery/plugin-sdk/v4/schema"
+	"github.com/cloudquery/plugin-sdk/v4/types"
+	"github.com/stretchr/testify/require"
+)
+
+const (
+	parquetString = "optional byte_array (String)"
+	parquetJSON   = parquetString
+	parquetBool   = "optional boolean"
+)
+
+func cloudflareCertificatePacksV11() *schema.Table {
+	return &schema.Table{Name: "cloudflare_certificate_packs", Columns: schema.ColumnList{
+		schema.CqIDColumn,
+		schema.CqParentIDColumn,
+		{Name: "account_id", Type: arrow.BinaryTypes.String},
+		{Name: "zone_id", Type: arrow.BinaryTypes.String},
+		{Name: "id", Type: arrow.BinaryTypes.String, PrimaryKey: true},
+		{Name: "primary_certificate", Type: arrow.BinaryTypes.String},
+		{Name: "certificate_authority", Type: arrow.BinaryTypes.String},
+		{Name: "type", Type: arrow.BinaryTypes.String},
+		{Name: "hosts", Type: arrow.ListOf(arrow.BinaryTypes.String)},
+		{Name: "status", Type: arrow.BinaryTypes.String},
+		{Name: "certificates", Type: types.ExtensionTypes.JSON},
+		{Name: "created_on", Type: arrow.BinaryTypes.String},
+		{Name: "validity_days", Type: arrow.PrimitiveTypes.Float64},
+		{Name: "validation_method", Type: arrow.BinaryTypes.String},
+	}}
+}
+
+func cloudflareCertificatePacksV12() *schema.Table {
+	return &schema.Table{Name: "cloudflare_certificate_packs", Columns: schema.ColumnList{
+		schema.CqIDColumn,
+		schema.CqParentIDColumn,
+		{Name: "account_id", Type: arrow.BinaryTypes.String},
+		{Name: "zone_id", Type: arrow.BinaryTypes.String},
+		{Name: "id", Type: arrow.BinaryTypes.String, PrimaryKey: true},
+		{Name: "validity_days", Type: arrow.PrimitiveTypes.Float64},
+		{Name: "certificates", Type: types.ExtensionTypes.JSON},
+		{Name: "hosts", Type: arrow.ListOf(arrow.BinaryTypes.String)},
+		{Name: "status", Type: arrow.BinaryTypes.String},
+		{Name: "type", Type: arrow.BinaryTypes.String},
+		{Name: "certificate_authority", Type: arrow.BinaryTypes.String},
+		{Name: "cloudflare_branding", Type: arrow.FixedWidthTypes.Boolean},
+		{Name: "primary_certificate", Type: arrow.BinaryTypes.String},
+		{Name: "validation_errors", Type: types.ExtensionTypes.JSON},
+		{Name: "validation_method", Type: arrow.BinaryTypes.String},
+		{Name: "validation_records", Type: types.ExtensionTypes.JSON},
+	}}
+}
+
+func TestParquetSchemaMatchesWrittenFile(t *testing.T) {
+	table := schema.TestTable("test", schema.TestSourceOptions{})
+	cl, err := NewClient()
+	require.NoError(t, err)
+
+	var b bytes.Buffer
+	h, err := cl.WriteHeader(&b, table)
+	require.NoError(t, err)
+	require.NoError(t, h.WriteFooter())
+
+	reader, err := file.NewParquetReader(bytes.NewReader(b.Bytes()))
+	require.NoError(t, err)
+	defer reader.Close()
+
+	generated, err := cl.ParquetSchema(table)
+	require.NoError(t, err)
+	require.True(t, reader.MetaData().Schema.Equals(generated))
+}
+
+func TestAssessTable(t *testing.T) {
+	stringList := arrow.ListOf(arrow.BinaryTypes.String)
+	tagsTable := func(dataType arrow.DataType) *schema.Table {
+		return &schema.Table{Name: "datadog_monitors", Columns: schema.ColumnList{
+			{Name: "id", Type: arrow.PrimitiveTypes.Int64, PrimaryKey: true, NotNull: true},
+			{Name: "tags", Type: dataType},
+		}}
+	}
+	tests := []struct {
+		name string
+		pair plugin.TablePair
+		want plugin.TableFinding
+	}{
+		{
+			name: "cloudflare certificate packs v11.4.0 to v12.0.0",
+			pair: plugin.TablePair{Old: cloudflareCertificatePacksV11(), New: cloudflareCertificatePacksV12()},
+			want: plugin.TableFinding{
+				TableName: "cloudflare_certificate_packs",
+				Category:  plugin.AssessCategoryFileSchemaChanged,
+				Columns: []plugin.ColumnFinding{
+					{ColumnName: "created_on", Category: plugin.AssessCategoryFileSchemaChanged, OldType: parquetString},
+					{ColumnName: "cloudflare_branding", Category: plugin.AssessCategoryFileSchemaChanged, NewType: parquetBool},
+					{ColumnName: "validation_errors", Category: plugin.AssessCategoryFileSchemaChanged, NewType: parquetJSON},
+					{ColumnName: "validation_records", Category: plugin.AssessCategoryFileSchemaChanged, NewType: parquetJSON},
+				},
+			},
+		},
+		{
+			name: "list of strings to json",
+			pair: plugin.TablePair{Old: tagsTable(stringList), New: tagsTable(types.ExtensionTypes.JSON)},
+			want: plugin.TableFinding{
+				TableName: "datadog_monitors",
+				Category:  plugin.AssessCategoryFileSchemaChanged,
+				Columns: []plugin.ColumnFinding{{
+					ColumnName: "tags",
+					Category:   plugin.AssessCategoryFileSchemaChanged,
+					OldType:    "optional group (List) {list: repeated group {element: optional byte_array (String)}}",
+					NewType:    parquetJSON,
+				}},
+			},
+		},
+		{
+			name: "types that write the same parquet type",
+			pair: plugin.TablePair{Old: tagsTable(arrow.BinaryTypes.LargeString), New: tagsTable(types.ExtensionTypes.UUID)},
+			want: plugin.TableFinding{TableName: "datadog_monitors", Category: plugin.AssessCategoryNoChange},
+		},
+		{
+			name: "timestamp precision change",
+			pair: plugin.TablePair{Old: tagsTable(arrow.FixedWidthTypes.Timestamp_ms), New: tagsTable(arrow.FixedWidthTypes.Timestamp_us)},
+			want: plugin.TableFinding{
+				TableName: "datadog_monitors",
+				Category:  plugin.AssessCategoryFileSchemaChanged,
+				Columns: []plugin.ColumnFinding{{
+					ColumnName: "tags",
+					Category:   plugin.AssessCategoryFileSchemaChanged,
+					OldType:    "optional int64 (Timestamp(isAdjustedToUTC=true, timeUnit=milliseconds, is_from_converted_type=false, force_set_converted_type=true))",
+					NewType:    "optional int64 (Timestamp(isAdjustedToUTC=true, timeUnit=microseconds, is_from_converted_type=false, force_set_converted_type=true))",
+				}},
+			},
+		},
+		{
+			name: "same table",
+			pair: plugin.TablePair{Old: cloudflareCertificatePacksV12(), New: cloudflareCertificatePacksV12()},
+			want: plugin.TableFinding{TableName: "cloudflare_certificate_packs", Category: plugin.AssessCategoryNoChange},
+		},
+		{
+			name: "added table",
+			pair: plugin.TablePair{New: tagsTable(stringList)},
+			want: plugin.TableFinding{
+				TableName: "datadog_monitors",
+				Category:  plugin.AssessCategoryFileSchemaChanged,
+				Columns: []plugin.ColumnFinding{
+					{ColumnName: "id", Category: plugin.AssessCategoryFileSchemaChanged, NewType: "required int64 (Int(bitWidth=64, isSigned=true))"},
+					{ColumnName: "tags", Category: plugin.AssessCategoryFileSchemaChanged, NewType: "optional group (List) {list: repeated group {element: optional byte_array (String)}}"},
+				},
+			},
+		},
+		{
+			name: "removed table",
+			pair: plugin.TablePair{Old: tagsTable(stringList)},
+			want: plugin.TableFinding{TableName: "datadog_monitors", Category: plugin.AssessCategoryTableRemoved},
+		},
+	}
+	cl, err := NewClient()
+	require.NoError(t, err)
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			got, err := cl.AssessTable(tc.pair)
+			require.NoError(t, err)
+			require.Equal(t, tc.want, got)
+		})
+	}
+}

parquet/write.go

@@ -21,16 +21,18 @@ type Handle struct {
 
 var _ ftypes.Handle = (*Handle)(nil)
 
-func (c *Client) WriteHeader(w io.Writer, t *schema.Table) (ftypes.Handle, error) {
-	props := parquet.NewWriterProperties(
+func (c *Client) writerProps() *parquet.WriterProperties {
+	return parquet.NewWriterProperties(
 		parquet.WithMaxRowGroupLength(c.spec.GetMaxRowGroupLength()),
 		parquet.WithCompression(compress.Codecs.Snappy),
 		parquet.WithVersion(c.spec.GetVersion()),
 		parquet.WithRootRepetition(c.spec.GetRootRepetition()),
 	)
-	arrprops := pqarrow.DefaultWriterProps()
+}
+
+func (c *Client) WriteHeader(w io.Writer, t *schema.Table) (ftypes.Handle, error) {
 	newSchema := convertSchema(t.ToArrowSchema())
-	fw, err := pqarrow.NewFileWriter(newSchema, &nopCloseWriter{Writer: w}, props, arrprops)
+	fw, err := pqarrow.NewFileWriter(newSchema, &nopCloseWriter{Writer: w}, c.writerProps(), pqarrow.DefaultWriterProps())
 	if err != nil {
 		return nil, err
 	}