From cc3f17981c76924983a69712a6bff8c93abb7034 Mon Sep 17 00:00:00 2001 From: Bruin Agent Date: Tue, 7 Jul 2026 14:18:23 +0000 Subject: [PATCH] Add Athena table summaries for data diff --- pkg/athena/db.go | 518 ++++++++++++++++++++++++++++++++++++++++- pkg/athena/db_test.go | 192 +++++++++++++++ pkg/diff/types.go | 44 ++++ pkg/diff/types_test.go | 54 +++++ 4 files changed, 803 insertions(+), 5 deletions(-) diff --git a/pkg/athena/db.go b/pkg/athena/db.go index 4a1d40c1b8..1589c18452 100644 --- a/pkg/athena/db.go +++ b/pkg/athena/db.go @@ -4,10 +4,12 @@ import ( "context" "fmt" "sort" + "strconv" "strings" "sync" "github.com/bruin-data/bruin/pkg/ansisql" + "github.com/bruin-data/bruin/pkg/diff" "github.com/bruin-data/bruin/pkg/query" "github.com/jmoiron/sqlx" "github.com/pkg/errors" @@ -15,15 +17,17 @@ import ( ) type DB struct { - conn *sqlx.DB - config *Config - mutex sync.Mutex + conn *sqlx.DB + config *Config + mutex sync.Mutex + typeMapper *diff.DatabaseTypeMapper } func NewDB(c *Config) *DB { return &DB{ - config: c, - mutex: sync.Mutex{}, + config: c, + mutex: sync.Mutex{}, + typeMapper: diff.NewAthenaTypeMapper(), } } @@ -393,6 +397,510 @@ ORDER BY table_schema, table_name; return summary, nil } +func (db *DB) GetTableSummary(ctx context.Context, tableName string, schemaOnly bool) (*diff.TableSummaryResult, error) { + schemaName, tableNameOnly, qualifiedTableName, err := db.parseTableName(tableName) + if err != nil { + return nil, err + } + + var rowCount int64 + if !schemaOnly { + rowCount, err = db.fetchRowCount(ctx, qualifiedTableName, tableName) + if err != nil { + return nil, err + } + } + + schemaResult, err := db.Select(ctx, &query.Query{Query: buildAthenaSchemaQuery(schemaName, tableNameOnly)}) + if err != nil { + return nil, fmt.Errorf("failed to execute schema query for table '%s': %w", tableName, err) + } + + typeMapper := db.typeMapper + if typeMapper == nil { + typeMapper = diff.NewAthenaTypeMapper() + } + + columns := make([]*diff.Column, 0, len(schemaResult)) + for _, row := range schemaResult { + if len(row) < 3 { + continue + } + + columnName, ok := row[0].(string) + if !ok { + continue + } + + dataType, ok := row[1].(string) + if !ok { + continue + } + + isNullableStr, ok := row[2].(string) + if !ok { + continue + } + + normalizedType := typeMapper.MapType(dataType) + nullable := strings.EqualFold(isNullableStr, "YES") + + var stats diff.ColumnStatistics + if schemaOnly { + stats = nil + } else { + switch normalizedType { + case diff.CommonTypeNumeric: + stats, err = db.fetchNumericalStats(ctx, qualifiedTableName, columnName) + if err != nil { + return nil, fmt.Errorf("failed to fetch numerical stats for column '%s': %w", columnName, err) + } + case diff.CommonTypeString: + stats, err = db.fetchStringStats(ctx, qualifiedTableName, columnName) + if err != nil { + return nil, fmt.Errorf("failed to fetch string stats for column '%s': %w", columnName, err) + } + case diff.CommonTypeBoolean: + stats, err = db.fetchBooleanStats(ctx, qualifiedTableName, columnName) + if err != nil { + return nil, fmt.Errorf("failed to fetch boolean stats for column '%s': %w", columnName, err) + } + case diff.CommonTypeDateTime: + stats, err = db.fetchDateTimeStats(ctx, qualifiedTableName, columnName) + if err != nil { + return nil, fmt.Errorf("failed to fetch datetime stats for column '%s': %w", columnName, err) + } + case diff.CommonTypeJSON: + stats, err = db.fetchJSONStats(ctx, qualifiedTableName, columnName) + if err != nil { + return nil, fmt.Errorf("failed to fetch JSON stats for column '%s': %w", columnName, err) + } + case diff.CommonTypeBinary, diff.CommonTypeUnknown: + stats = &diff.UnknownStatistics{} + } + } + + columns = append(columns, &diff.Column{ + Name: columnName, + Type: dataType, + NormalizedType: normalizedType, + Nullable: nullable, + PrimaryKey: false, + Unique: false, + Stats: stats, + }) + } + + return &diff.TableSummaryResult{ + RowCount: rowCount, + Table: &diff.Table{ + Name: tableName, + Columns: columns, + }, + }, nil +} + +func (db *DB) parseTableName(tableName string) (string, string, string, error) { + tableComponents := strings.Split(tableName, ".") + for _, component := range tableComponents { + if component == "" { + return "", "", "", fmt.Errorf("table name must be in table or schema.table format, '%s' given", tableName) + } + } + + var schemaName string + var tableNameOnly string + + switch len(tableComponents) { + case 1: + if db.config == nil || db.config.Database == "" { + return "", "", "", fmt.Errorf("database must be configured when table name is not schema-qualified: %s", tableName) + } + schemaName = db.config.Database + tableNameOnly = tableComponents[0] + case 2: + schemaName = tableComponents[0] + tableNameOnly = tableComponents[1] + default: + return "", "", "", fmt.Errorf("table name must be in table or schema.table format, '%s' given", tableName) + } + + return schemaName, tableNameOnly, quoteAthenaQualifiedTableName(schemaName, tableNameOnly), nil +} + +func buildAthenaSchemaQuery(schemaName, tableName string) string { + return fmt.Sprintf(` +SELECT + column_name, + data_type, + is_nullable +FROM information_schema.columns +WHERE table_schema = '%s' AND table_name = '%s' +ORDER BY ordinal_position; +`, escapeAthenaStringLiteral(schemaName), escapeAthenaStringLiteral(tableName)) +} + +func (db *DB) fetchRowCount(ctx context.Context, qualifiedTableName, originalTableName string) (int64, error) { + countQuery := fmt.Sprintf("SELECT COUNT(*) as row_count FROM %s", qualifiedTableName) + countResult, err := db.Select(ctx, &query.Query{Query: countQuery}) + if err != nil { + return 0, fmt.Errorf("failed to execute count query for table '%s': %w", originalTableName, err) + } + + if len(countResult) == 0 || len(countResult[0]) == 0 { + return 0, fmt.Errorf("count query returned no rows for table '%s'", originalTableName) + } + + rowCount, err := athenaInt64Value(countResult[0][0]) + if err != nil { + return 0, fmt.Errorf("failed to parse row count for table '%s': %w", originalTableName, err) + } + + return rowCount, nil +} + +func (db *DB) fetchNumericalStats(ctx context.Context, qualifiedTableName, columnName string) (*diff.NumericalStatistics, error) { + quotedColumn := quoteAthenaIdentifier(columnName) + statsQuery := fmt.Sprintf(` +SELECT + MIN(TRY_CAST(%s AS DOUBLE)) as min_val, + MAX(TRY_CAST(%s AS DOUBLE)) as max_val, + AVG(TRY_CAST(%s AS DOUBLE)) as avg_val, + SUM(TRY_CAST(%s AS DOUBLE)) as sum_val, + COUNT(%s) as count_val, + COUNT(*) - COUNT(%s) as null_count, + STDDEV(TRY_CAST(%s AS DOUBLE)) as stddev_val +FROM %s +`, quotedColumn, quotedColumn, quotedColumn, quotedColumn, quotedColumn, quotedColumn, quotedColumn, qualifiedTableName) + + result, err := db.Select(ctx, &query.Query{Query: statsQuery}) + if err != nil { + return nil, fmt.Errorf("failed to fetch numerical stats for column '%s': %w", columnName, err) + } + if len(result) == 0 || len(result[0]) < 7 { + return nil, fmt.Errorf("insufficient statistical data returned for column '%s'", columnName) + } + + row := result[0] + stats := &diff.NumericalStatistics{} + + if stats.Min, err = athenaOptionalFloat64Value(row[0]); err != nil { + return nil, fmt.Errorf("failed to parse min value for column '%s': %w", columnName, err) + } + if stats.Max, err = athenaOptionalFloat64Value(row[1]); err != nil { + return nil, fmt.Errorf("failed to parse max value for column '%s': %w", columnName, err) + } + if stats.Avg, err = athenaOptionalFloat64Value(row[2]); err != nil { + return nil, fmt.Errorf("failed to parse avg value for column '%s': %w", columnName, err) + } + if stats.Sum, err = athenaOptionalFloat64Value(row[3]); err != nil { + return nil, fmt.Errorf("failed to parse sum value for column '%s': %w", columnName, err) + } + if stats.Count, err = athenaInt64Value(row[4]); err != nil { + return nil, fmt.Errorf("failed to parse count value for column '%s': %w", columnName, err) + } + if stats.NullCount, err = athenaInt64Value(row[5]); err != nil { + return nil, fmt.Errorf("failed to parse null count for column '%s': %w", columnName, err) + } + if stats.StdDev, err = athenaOptionalFloat64Value(row[6]); err != nil { + return nil, fmt.Errorf("failed to parse stddev value for column '%s': %w", columnName, err) + } + + return stats, nil +} + +func (db *DB) fetchStringStats(ctx context.Context, qualifiedTableName, columnName string) (*diff.StringStatistics, error) { + quotedColumn := quoteAthenaIdentifier(columnName) + statsQuery := fmt.Sprintf(` +SELECT + MIN(LENGTH(CAST(%s AS VARCHAR))) as min_len, + MAX(LENGTH(CAST(%s AS VARCHAR))) as max_len, + AVG(LENGTH(CAST(%s AS VARCHAR))) as avg_len, + COUNT(DISTINCT %s) as distinct_count, + COUNT(*) as total_count, + COUNT(*) - COUNT(%s) as null_count, + SUM(CASE WHEN CAST(%s AS VARCHAR) = '' THEN 1 ELSE 0 END) as empty_count +FROM %s +`, quotedColumn, quotedColumn, quotedColumn, quotedColumn, quotedColumn, quotedColumn, qualifiedTableName) + + result, err := db.Select(ctx, &query.Query{Query: statsQuery}) + if err != nil { + return nil, fmt.Errorf("failed to fetch string stats for column '%s': %w", columnName, err) + } + if len(result) == 0 || len(result[0]) < 7 { + return nil, fmt.Errorf("insufficient statistical data returned for column '%s'", columnName) + } + + row := result[0] + stats := &diff.StringStatistics{} + + if stats.MinLength, err = athenaIntValue(row[0]); err != nil { + return nil, fmt.Errorf("failed to parse min length for column '%s': %w", columnName, err) + } + if stats.MaxLength, err = athenaIntValue(row[1]); err != nil { + return nil, fmt.Errorf("failed to parse max length for column '%s': %w", columnName, err) + } + if stats.AvgLength, err = athenaFloat64Value(row[2]); err != nil { + return nil, fmt.Errorf("failed to parse avg length for column '%s': %w", columnName, err) + } + if stats.DistinctCount, err = athenaInt64Value(row[3]); err != nil { + return nil, fmt.Errorf("failed to parse distinct count for column '%s': %w", columnName, err) + } + if stats.Count, err = athenaInt64Value(row[4]); err != nil { + return nil, fmt.Errorf("failed to parse count value for column '%s': %w", columnName, err) + } + if stats.NullCount, err = athenaInt64Value(row[5]); err != nil { + return nil, fmt.Errorf("failed to parse null count for column '%s': %w", columnName, err) + } + if stats.EmptyCount, err = athenaInt64Value(row[6]); err != nil { + return nil, fmt.Errorf("failed to parse empty count for column '%s': %w", columnName, err) + } + + return stats, nil +} + +func (db *DB) fetchBooleanStats(ctx context.Context, qualifiedTableName, columnName string) (*diff.BooleanStatistics, error) { + quotedColumn := quoteAthenaIdentifier(columnName) + statsQuery := fmt.Sprintf(` +SELECT + SUM(CASE WHEN %s = true THEN 1 ELSE 0 END) as true_count, + SUM(CASE WHEN %s = false THEN 1 ELSE 0 END) as false_count, + COUNT(*) as total_count, + COUNT(*) - COUNT(%s) as null_count +FROM %s +`, quotedColumn, quotedColumn, quotedColumn, qualifiedTableName) + + result, err := db.Select(ctx, &query.Query{Query: statsQuery}) + if err != nil { + return nil, fmt.Errorf("failed to fetch boolean stats for column '%s': %w", columnName, err) + } + if len(result) == 0 || len(result[0]) < 4 { + return nil, fmt.Errorf("insufficient statistical data returned for column '%s'", columnName) + } + + row := result[0] + stats := &diff.BooleanStatistics{} + + if stats.TrueCount, err = athenaInt64Value(row[0]); err != nil { + return nil, fmt.Errorf("failed to parse true count for column '%s': %w", columnName, err) + } + if stats.FalseCount, err = athenaInt64Value(row[1]); err != nil { + return nil, fmt.Errorf("failed to parse false count for column '%s': %w", columnName, err) + } + if stats.Count, err = athenaInt64Value(row[2]); err != nil { + return nil, fmt.Errorf("failed to parse count value for column '%s': %w", columnName, err) + } + if stats.NullCount, err = athenaInt64Value(row[3]); err != nil { + return nil, fmt.Errorf("failed to parse null count for column '%s': %w", columnName, err) + } + + return stats, nil +} + +func (db *DB) fetchDateTimeStats(ctx context.Context, qualifiedTableName, columnName string) (*diff.DateTimeStatistics, error) { + quotedColumn := quoteAthenaIdentifier(columnName) + statsQuery := fmt.Sprintf(` +SELECT + CAST(MIN(%s) AS VARCHAR) as min_date, + CAST(MAX(%s) AS VARCHAR) as max_date, + COUNT(DISTINCT %s) as unique_count, + COUNT(*) as count_val, + COUNT(*) - COUNT(%s) as null_count +FROM %s +`, quotedColumn, quotedColumn, quotedColumn, quotedColumn, qualifiedTableName) + + result, err := db.Select(ctx, &query.Query{Query: statsQuery}) + if err != nil { + return nil, fmt.Errorf("failed to fetch datetime stats for column '%s': %w", columnName, err) + } + if len(result) == 0 || len(result[0]) < 5 { + return nil, fmt.Errorf("insufficient statistical data returned for column '%s'", columnName) + } + + row := result[0] + stats := &diff.DateTimeStatistics{} + + if row[0] != nil { + if parsedTime, parseErr := diff.ParseDateTime(row[0]); parseErr == nil { + stats.EarliestDate = parsedTime + } + } + if row[1] != nil { + if parsedTime, parseErr := diff.ParseDateTime(row[1]); parseErr == nil { + stats.LatestDate = parsedTime + } + } + if stats.UniqueCount, err = athenaInt64Value(row[2]); err != nil { + return nil, fmt.Errorf("failed to parse unique count for column '%s': %w", columnName, err) + } + if stats.Count, err = athenaInt64Value(row[3]); err != nil { + return nil, fmt.Errorf("failed to parse count value for column '%s': %w", columnName, err) + } + if stats.NullCount, err = athenaInt64Value(row[4]); err != nil { + return nil, fmt.Errorf("failed to parse null count for column '%s': %w", columnName, err) + } + + return stats, nil +} + +func (db *DB) fetchJSONStats(ctx context.Context, qualifiedTableName, columnName string) (*diff.JSONStatistics, error) { + quotedColumn := quoteAthenaIdentifier(columnName) + statsQuery := fmt.Sprintf(` +SELECT + COUNT(*) as count_val, + COUNT(*) - COUNT(%s) as null_count +FROM %s +`, quotedColumn, qualifiedTableName) + + result, err := db.Select(ctx, &query.Query{Query: statsQuery}) + if err != nil { + return nil, fmt.Errorf("failed to fetch JSON stats for column '%s': %w", columnName, err) + } + if len(result) == 0 || len(result[0]) < 2 { + return nil, fmt.Errorf("insufficient statistical data returned for column '%s'", columnName) + } + + row := result[0] + stats := &diff.JSONStatistics{} + + if stats.Count, err = athenaInt64Value(row[0]); err != nil { + return nil, fmt.Errorf("failed to parse count value for column '%s': %w", columnName, err) + } + if stats.NullCount, err = athenaInt64Value(row[1]); err != nil { + return nil, fmt.Errorf("failed to parse null count for column '%s': %w", columnName, err) + } + + return stats, nil +} + +func athenaIntValue(value interface{}) (int, error) { + int64Value, err := athenaInt64Value(value) + if err != nil { + return 0, err + } + + return int(int64Value), nil +} + +func athenaInt64Value(value interface{}) (int64, error) { + switch val := value.(type) { + case nil: + return 0, nil + case int64: + return val, nil + case int: + return int64(val), nil + case int32: + return int64(val), nil + case int16: + return int64(val), nil + case int8: + return int64(val), nil + case uint64: + return int64(val), nil + case uint: + return int64(val), nil + case uint32: + return int64(val), nil + case uint16: + return int64(val), nil + case uint8: + return int64(val), nil + case float64: + return int64(val), nil + case float32: + return int64(val), nil + case []byte: + return athenaInt64Value(string(val)) + case string: + trimmed := strings.TrimSpace(val) + if trimmed == "" { + return 0, nil + } + parsed, err := strconv.ParseInt(trimmed, 10, 64) + if err == nil { + return parsed, nil + } + parsedFloat, floatErr := strconv.ParseFloat(trimmed, 64) + if floatErr != nil { + return 0, err + } + return int64(parsedFloat), nil + default: + return 0, fmt.Errorf("unexpected numeric value type %T with value %v", val, val) + } +} + +func athenaFloat64Value(value interface{}) (float64, error) { + switch val := value.(type) { + case nil: + return 0, nil + case float64: + return val, nil + case float32: + return float64(val), nil + case int64: + return float64(val), nil + case int: + return float64(val), nil + case int32: + return float64(val), nil + case int16: + return float64(val), nil + case int8: + return float64(val), nil + case uint64: + return float64(val), nil + case uint: + return float64(val), nil + case uint32: + return float64(val), nil + case uint16: + return float64(val), nil + case uint8: + return float64(val), nil + case []byte: + return athenaFloat64Value(string(val)) + case string: + trimmed := strings.TrimSpace(val) + if trimmed == "" { + return 0, nil + } + parsed, err := strconv.ParseFloat(trimmed, 64) + if err != nil { + return 0, err + } + return parsed, nil + default: + return 0, fmt.Errorf("unexpected numeric value type %T with value %v", val, val) + } +} + +func athenaOptionalFloat64Value(value interface{}) (*float64, error) { + if value == nil { + return nil, nil + } + + parsed, err := athenaFloat64Value(value) + if err != nil { + return nil, err + } + + return &parsed, nil +} + +func quoteAthenaQualifiedTableName(schemaName, tableName string) string { + return quoteAthenaIdentifier(schemaName) + "." + quoteAthenaIdentifier(tableName) +} + +func quoteAthenaIdentifier(identifier string) string { + return `"` + strings.ReplaceAll(identifier, `"`, `""`) + `"` +} + +func escapeAthenaStringLiteral(value string) string { + return strings.ReplaceAll(value, "'", "''") +} + func (db *DB) BuildTableExistsQuery(tableName string) (string, error) { tableComponents := strings.Split(tableName, ".") diff --git a/pkg/athena/db_test.go b/pkg/athena/db_test.go index 20a6a5e04d..fe8371f2e8 100644 --- a/pkg/athena/db_test.go +++ b/pkg/athena/db_test.go @@ -5,6 +5,7 @@ import ( "testing" "github.com/DATA-DOG/go-sqlmock" + "github.com/bruin-data/bruin/pkg/diff" "github.com/bruin-data/bruin/pkg/query" "github.com/jmoiron/sqlx" "github.com/stretchr/testify/assert" @@ -264,6 +265,197 @@ func TestDB_SelectWithSchema(t *testing.T) { } } +func TestDB_GetTableSummarySchemaOnly(t *testing.T) { + t.Parallel() + + mockDB, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual)) + require.NoError(t, err) + defer mockDB.Close() + + sqlxDB := sqlx.NewDb(mockDB, "sqlmock") + mock.ExpectQuery(buildAthenaSchemaQuery("analytics", "orders")). + WillReturnRows( + sqlmock.NewRows([]string{"column_name", "data_type", "is_nullable"}). + AddRow("id", "bigint", "NO"). + AddRow("status", "varchar", "YES"). + AddRow("metadata", "json", "YES"), + ) + + db := DB{ + conn: sqlxDB, + config: &Config{Database: "analytics"}, + typeMapper: diff.NewAthenaTypeMapper(), + } + + got, err := db.GetTableSummary(t.Context(), "orders", true) + require.NoError(t, err) + + require.Equal(t, int64(0), got.RowCount) + require.Equal(t, "orders", got.Table.Name) + require.Len(t, got.Table.Columns, 3) + + assert.Equal(t, "id", got.Table.Columns[0].Name) + assert.Equal(t, "bigint", got.Table.Columns[0].Type) + assert.Equal(t, diff.CommonTypeNumeric, got.Table.Columns[0].NormalizedType) + assert.False(t, got.Table.Columns[0].Nullable) + assert.Nil(t, got.Table.Columns[0].Stats) + + assert.Equal(t, "status", got.Table.Columns[1].Name) + assert.Equal(t, diff.CommonTypeString, got.Table.Columns[1].NormalizedType) + assert.True(t, got.Table.Columns[1].Nullable) + assert.Nil(t, got.Table.Columns[1].Stats) + + assert.Equal(t, "metadata", got.Table.Columns[2].Name) + assert.Equal(t, diff.CommonTypeJSON, got.Table.Columns[2].NormalizedType) + assert.Nil(t, got.Table.Columns[2].Stats) + + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestDB_GetTableSummaryFull(t *testing.T) { + t.Parallel() + + mockDB, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual)) + require.NoError(t, err) + defer mockDB.Close() + + sqlxDB := sqlx.NewDb(mockDB, "sqlmock") + mock.ExpectQuery(`SELECT COUNT(*) as row_count FROM "analytics"."orders"`). + WillReturnRows(sqlmock.NewRows([]string{"row_count"}).AddRow(4)) + mock.ExpectQuery(buildAthenaSchemaQuery("analytics", "orders")). + WillReturnRows( + sqlmock.NewRows([]string{"column_name", "data_type", "is_nullable"}). + AddRow("amount", "decimal(10,2)", "YES"). + AddRow("customer_name", "varchar", "YES"). + AddRow("paid", "boolean", "YES"). + AddRow("created_at", "timestamp", "NO"). + AddRow("payload", "json", "YES"). + AddRow("raw_bytes", "varbinary", "YES"), + ) + mock.ExpectQuery(` +SELECT + MIN(TRY_CAST("amount" AS DOUBLE)) as min_val, + MAX(TRY_CAST("amount" AS DOUBLE)) as max_val, + AVG(TRY_CAST("amount" AS DOUBLE)) as avg_val, + SUM(TRY_CAST("amount" AS DOUBLE)) as sum_val, + COUNT("amount") as count_val, + COUNT(*) - COUNT("amount") as null_count, + STDDEV(TRY_CAST("amount" AS DOUBLE)) as stddev_val +FROM "analytics"."orders" +`).WillReturnRows(sqlmock.NewRows([]string{"min_val", "max_val", "avg_val", "sum_val", "count_val", "null_count", "stddev_val"}). + AddRow(1.5, 9.5, 5.5, 22.0, 4, 1, 3.1)) + mock.ExpectQuery(` +SELECT + MIN(LENGTH(CAST("customer_name" AS VARCHAR))) as min_len, + MAX(LENGTH(CAST("customer_name" AS VARCHAR))) as max_len, + AVG(LENGTH(CAST("customer_name" AS VARCHAR))) as avg_len, + COUNT(DISTINCT "customer_name") as distinct_count, + COUNT(*) as total_count, + COUNT(*) - COUNT("customer_name") as null_count, + SUM(CASE WHEN CAST("customer_name" AS VARCHAR) = '' THEN 1 ELSE 0 END) as empty_count +FROM "analytics"."orders" +`).WillReturnRows(sqlmock.NewRows([]string{"min_len", "max_len", "avg_len", "distinct_count", "total_count", "null_count", "empty_count"}). + AddRow(3, 12, 7.5, 3, 4, 1, 1)) + mock.ExpectQuery(` +SELECT + SUM(CASE WHEN "paid" = true THEN 1 ELSE 0 END) as true_count, + SUM(CASE WHEN "paid" = false THEN 1 ELSE 0 END) as false_count, + COUNT(*) as total_count, + COUNT(*) - COUNT("paid") as null_count +FROM "analytics"."orders" +`).WillReturnRows(sqlmock.NewRows([]string{"true_count", "false_count", "total_count", "null_count"}). + AddRow(2, 1, 4, 1)) + mock.ExpectQuery(` +SELECT + CAST(MIN("created_at") AS VARCHAR) as min_date, + CAST(MAX("created_at") AS VARCHAR) as max_date, + COUNT(DISTINCT "created_at") as unique_count, + COUNT(*) as count_val, + COUNT(*) - COUNT("created_at") as null_count +FROM "analytics"."orders" +`).WillReturnRows(sqlmock.NewRows([]string{"min_date", "max_date", "unique_count", "count_val", "null_count"}). + AddRow("2024-01-01 00:00:00", "2024-01-03 12:00:00", 3, 4, 1)) + mock.ExpectQuery(` +SELECT + COUNT(*) as count_val, + COUNT(*) - COUNT("payload") as null_count +FROM "analytics"."orders" +`).WillReturnRows(sqlmock.NewRows([]string{"count_val", "null_count"}).AddRow(4, 1)) + + db := DB{ + conn: sqlxDB, + config: &Config{Database: "unused"}, + typeMapper: diff.NewAthenaTypeMapper(), + } + + got, err := db.GetTableSummary(t.Context(), "analytics.orders", false) + require.NoError(t, err) + + require.Equal(t, int64(4), got.RowCount) + require.Equal(t, "analytics.orders", got.Table.Name) + require.Len(t, got.Table.Columns, 6) + + amountStats, ok := got.Table.Columns[0].Stats.(*diff.NumericalStatistics) + require.True(t, ok) + require.NotNil(t, amountStats.Min) + require.NotNil(t, amountStats.Max) + require.NotNil(t, amountStats.Avg) + require.NotNil(t, amountStats.Sum) + require.NotNil(t, amountStats.StdDev) + assert.InDelta(t, 1.5, *amountStats.Min, 0.0001) + assert.InDelta(t, 9.5, *amountStats.Max, 0.0001) + assert.InDelta(t, 5.5, *amountStats.Avg, 0.0001) + assert.InDelta(t, 22.0, *amountStats.Sum, 0.0001) + assert.Equal(t, int64(4), amountStats.Count) + assert.Equal(t, int64(1), amountStats.NullCount) + assert.InDelta(t, 3.1, *amountStats.StdDev, 0.0001) + + stringStats, ok := got.Table.Columns[1].Stats.(*diff.StringStatistics) + require.True(t, ok) + assert.Equal(t, 3, stringStats.MinLength) + assert.Equal(t, 12, stringStats.MaxLength) + assert.InDelta(t, 7.5, stringStats.AvgLength, 0.0001) + assert.Equal(t, int64(3), stringStats.DistinctCount) + assert.Equal(t, int64(4), stringStats.Count) + assert.Equal(t, int64(1), stringStats.NullCount) + assert.Equal(t, int64(1), stringStats.EmptyCount) + + booleanStats, ok := got.Table.Columns[2].Stats.(*diff.BooleanStatistics) + require.True(t, ok) + assert.Equal(t, int64(2), booleanStats.TrueCount) + assert.Equal(t, int64(1), booleanStats.FalseCount) + assert.Equal(t, int64(4), booleanStats.Count) + assert.Equal(t, int64(1), booleanStats.NullCount) + + dateTimeStats, ok := got.Table.Columns[3].Stats.(*diff.DateTimeStatistics) + require.True(t, ok) + require.NotNil(t, dateTimeStats.EarliestDate) + require.NotNil(t, dateTimeStats.LatestDate) + assert.Equal(t, int64(3), dateTimeStats.UniqueCount) + assert.Equal(t, int64(4), dateTimeStats.Count) + assert.Equal(t, int64(1), dateTimeStats.NullCount) + + jsonStats, ok := got.Table.Columns[4].Stats.(*diff.JSONStatistics) + require.True(t, ok) + assert.Equal(t, int64(4), jsonStats.Count) + assert.Equal(t, int64(1), jsonStats.NullCount) + + _, ok = got.Table.Columns[5].Stats.(*diff.UnknownStatistics) + require.True(t, ok) + + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestDB_GetTableSummaryInvalidTableName(t *testing.T) { + t.Parallel() + + db := DB{config: &Config{Database: "analytics"}} + + _, err := db.GetTableSummary(t.Context(), "a.b.c", true) + require.Error(t, err) + assert.Contains(t, err.Error(), "table name must be in table or schema.table format") +} + func TestDB_BuildTableExistsQuery(t *testing.T) { t.Parallel() tests := []struct { diff --git a/pkg/diff/types.go b/pkg/diff/types.go index b557892665..1cde4f6741 100644 --- a/pkg/diff/types.go +++ b/pkg/diff/types.go @@ -247,6 +247,50 @@ func NewBigQueryTypeMapper() *DatabaseTypeMapper { return mapper } +// NewAthenaTypeMapper provides AWS Athena/Presto-specific type mapping. +func NewAthenaTypeMapper() *DatabaseTypeMapper { + mapper := NewDatabaseTypeMapper() + + // Numeric types in Athena. + mapper.AddNumericTypes( + "tinyint", "smallint", "integer", "int", "bigint", + "real", "double", "double precision", + "decimal", "numeric", + ) + + // String types in Athena. + mapper.AddStringTypes( + "char", "varchar", "string", + "uuid", "ipaddress", + ) + + // Boolean types in Athena. + mapper.AddBooleanTypes( + "boolean", "bool", + ) + + // DateTime types in Athena. + mapper.AddDateTimeTypes( + "date", + "time", "time with time zone", + "timestamp", "timestamp with time zone", + "interval year to month", "interval day to second", + ) + + // Binary types in Athena. + mapper.AddBinaryTypes( + "varbinary", + ) + + // JSON type in Athena. Complex ARRAY/MAP/ROW types remain unknown because + // their value-level comparison semantics vary by nested shape. + mapper.AddJSONTypes( + "json", + ) + + return mapper +} + // NewPostgresTypeMapper provides PostgreSQL-specific type mapping. func NewPostgresTypeMapper() *DatabaseTypeMapper { mapper := NewDatabaseTypeMapper() diff --git a/pkg/diff/types_test.go b/pkg/diff/types_test.go index d7f5d08e1a..c302e7ccd7 100644 --- a/pkg/diff/types_test.go +++ b/pkg/diff/types_test.go @@ -256,6 +256,60 @@ func TestBigQueryTypeMapper(t *testing.T) { } } +func TestAthenaTypeMapper(t *testing.T) { + t.Parallel() + mapper := NewAthenaTypeMapper() + + tests := []struct { + inputType string + expectedResult CommonDataType + }{ + // Athena numeric types + {"tinyint", CommonTypeNumeric}, + {"smallint", CommonTypeNumeric}, + {"integer", CommonTypeNumeric}, + {"bigint", CommonTypeNumeric}, + {"real", CommonTypeNumeric}, + {"double", CommonTypeNumeric}, + {"decimal(10,2)", CommonTypeNumeric}, + + // Athena string types + {"varchar", CommonTypeString}, + {"varchar(255)", CommonTypeString}, + {"char", CommonTypeString}, + {"string", CommonTypeString}, + {"ipaddress", CommonTypeString}, + + // Athena boolean types + {"boolean", CommonTypeBoolean}, + + // Athena datetime types + {"date", CommonTypeDateTime}, + {"timestamp", CommonTypeDateTime}, + {"timestamp with time zone", CommonTypeDateTime}, + {"time", CommonTypeDateTime}, + + // Athena binary and JSON types + {"varbinary", CommonTypeBinary}, + {"json", CommonTypeJSON}, + + // Athena complex types are intentionally not normalized. + {"array(varchar)", CommonTypeUnknown}, + {"map(varchar,bigint)", CommonTypeUnknown}, + {"row(id bigint)", CommonTypeUnknown}, + } + + for _, tt := range tests { + t.Run("Athena_"+tt.inputType, func(t *testing.T) { + t.Parallel() + result := mapper.MapType(tt.inputType) + if result != tt.expectedResult { + t.Errorf("Athena mapper: MapType(%q) = %v, want %v", tt.inputType, result, tt.expectedResult) + } + }) + } +} + func TestDatabaseTypeMapper_IsNumeric(t *testing.T) { t.Parallel() mapper := createTestMapper()