diff --git a/agents/drivers/kingbase-go/kingbase_metadata.go b/agents/drivers/kingbase-go/kingbase_metadata.go index b1981e363..eef6bf2d2 100644 --- a/agents/drivers/kingbase-go/kingbase_metadata.go +++ b/agents/drivers/kingbase-go/kingbase_metadata.go @@ -62,6 +62,7 @@ type metadataListConstraints struct { type columnInfo struct { Name string `json:"name"` DataType string `json:"data_type"` + FullDataType string `json:"-"` IsNullable bool `json:"is_nullable"` ColumnDefault *string `json:"column_default"` IsPrimaryKey bool `json:"is_primary_key"` @@ -538,13 +539,25 @@ func isUndefinedColumn(err error, columnName string) bool { } func (s *server) informationSchemaColumns(schema, table string, primary map[string]bool) ([]columnInfo, error) { - query := `SELECT c.column_name, c.data_type, c.is_nullable, c.column_default, + result, err := s.queryInformationSchemaColumns(schema, table, primary, true) + if err != nil && isUndefinedColumn(err, "column_type") { + return s.queryInformationSchemaColumns(schema, table, primary, false) + } + return result, err +} + +func (s *server) queryInformationSchemaColumns(schema, table string, primary map[string]bool, includeFullDataType bool) ([]columnInfo, error) { + fullDataTypeExpression := "c.column_type" + if !includeFullDataType { + fullDataTypeExpression = "NULL AS column_type" + } + query := fmt.Sprintf(`SELECT c.column_name, c.data_type, %s, c.is_nullable, c.column_default, col_description(a.attrelid, a.attnum), c.numeric_precision, c.numeric_scale, c.character_maximum_length FROM information_schema.columns c LEFT JOIN sys_catalog.sys_namespace n ON n.nspname = c.table_schema LEFT JOIN sys_catalog.sys_class rel ON rel.relnamespace = n.oid AND rel.relname = c.table_name LEFT JOIN sys_catalog.sys_attribute a ON a.attrelid = rel.oid AND a.attname = c.column_name AND a.attnum > 0 AND NOT a.attisdropped - WHERE c.table_schema = ` + quoteLiteral(schema) + ` AND c.table_name = ` + quoteLiteral(table) + ` ORDER BY c.ordinal_position` + WHERE c.table_schema = %s AND c.table_name = %s ORDER BY c.ordinal_position`, fullDataTypeExpression, quoteLiteral(schema), quoteLiteral(table)) rows, err := s.metadataQuery(query) if err != nil { return nil, err @@ -553,15 +566,15 @@ func (s *server) informationSchemaColumns(schema, table string, primary map[stri result := []columnInfo{} for rows.Next() { var name, dataType, nullable string - var defaultValue, comment sql.NullString + var fullDataType, defaultValue, comment sql.NullString var precision, scale, length sql.NullInt64 - if err := rows.Scan(&name, &dataType, &nullable, &defaultValue, &comment, &precision, &scale, &length); err != nil { + if err := rows.Scan(&name, &dataType, &fullDataType, &nullable, &defaultValue, &comment, &precision, &scale, &length); err != nil { return nil, err } if parsed := boundedVarcharLength(dataType); parsed != nil && !length.Valid { length = sql.NullInt64{Int64: int64(*parsed), Valid: true} } - result = append(result, columnInfo{Name: name, DataType: dataType, IsNullable: strings.EqualFold(nullable, "YES"), ColumnDefault: nullStringPtr(defaultValue), IsPrimaryKey: primary[strings.ToLower(name)], Comment: nullStringPtr(comment), NumericPrecision: nullIntPtr(precision), NumericScale: nullIntPtr(scale), CharacterMaximumLength: nullIntPtr(length)}) + result = append(result, columnInfo{Name: name, DataType: dataType, FullDataType: fullDataType.String, IsNullable: strings.EqualFold(nullable, "YES"), ColumnDefault: nullStringPtr(defaultValue), IsPrimaryKey: primary[strings.ToLower(name)], Comment: nullStringPtr(comment), NumericPrecision: nullIntPtr(precision), NumericScale: nullIntPtr(scale), CharacterMaximumLength: nullIntPtr(length)}) } return result, rows.Err() } @@ -875,7 +888,7 @@ func ensureStatementTerminator(statement string) string { } func columnDDLDefinition(column columnInfo) string { - definition := quoteIdentifier(column.Name) + " " + column.DataType + definition := quoteIdentifier(column.Name) + " " + columnDDLDataType(column) if column.Extra != nil && *column.Extra != "" { // Identity clauses belong immediately after the data type in both // PostgreSQL-compatible and SQL Server-compatible Kingbase modes. @@ -890,6 +903,31 @@ func columnDDLDefinition(column columnInfo) string { return definition } +func columnDDLDataType(column columnInfo) string { + if fullDataType := strings.TrimSpace(column.FullDataType); fullDataType != "" { + return fullDataType + } + dataType := strings.TrimSpace(column.DataType) + if strings.Contains(dataType, "(") { + return dataType + } + normalized := strings.Join(strings.Fields(strings.ToLower(dataType)), " ") + switch normalized { + case "varchar", "character varying", "char", "character": + if column.CharacterMaximumLength != nil && *column.CharacterMaximumLength > 0 { + return fmt.Sprintf("%s(%d)", dataType, *column.CharacterMaximumLength) + } + case "numeric", "decimal": + if column.NumericPrecision != nil && *column.NumericPrecision > 0 { + if column.NumericScale != nil { + return fmt.Sprintf("%s(%d,%d)", dataType, *column.NumericPrecision, *column.NumericScale) + } + return fmt.Sprintf("%s(%d)", dataType, *column.NumericPrecision) + } + } + return dataType +} + func kingbaseIdentityClause(code string) *string { var clause string switch strings.ToLower(strings.TrimSpace(code)) { diff --git a/agents/drivers/kingbase-go/main_test.go b/agents/drivers/kingbase-go/main_test.go index 8c17b04f5..8efd3f663 100644 --- a/agents/drivers/kingbase-go/main_test.go +++ b/agents/drivers/kingbase-go/main_test.go @@ -204,8 +204,8 @@ func (connection *fallbackConn) QueryContext(_ context.Context, query string, _ } if strings.Contains(query, "FROM information_schema.columns c") { return &valueRows{ - columns: []string{"column_name", "data_type", "is_nullable", "column_default", "column_comment", "numeric_precision", "numeric_scale", "character_maximum_length"}, - rows: [][]driver.Value{{"id", "integer", "NO", nil, "primary key", int64(32), int64(0), nil}}, + columns: []string{"column_name", "data_type", "column_type", "is_nullable", "column_default", "column_comment", "numeric_precision", "numeric_scale", "character_maximum_length"}, + rows: [][]driver.Value{{"id", "integer", "integer", "NO", nil, "primary key", int64(32), int64(0), nil}}, }, nil } if connection.state.rejectAttidentity && strings.Contains(query, "a.attidentity") { @@ -1065,6 +1065,125 @@ func TestColumnDDLDefinitionPreservesCompatibilityExtras(t *testing.T) { } } +func TestColumnDDLDefinitionRestoresMySQLCompatibilityTypeModifiers(t *testing.T) { + length := 64 + precision := 12 + scale := 4 + tests := []struct { + name string + column columnInfo + want string + }{ + {name: "varchar length", column: columnInfo{Name: "label", DataType: "varchar", IsNullable: true, CharacterMaximumLength: &length}, want: `"label" varchar(64)`}, + {name: "numeric precision and scale", column: columnInfo{Name: "amount", DataType: "numeric", IsNullable: true, NumericPrecision: &precision, NumericScale: &scale}, want: `"amount" numeric(12,4)`}, + {name: "existing modifier", column: columnInfo{Name: "code", DataType: "VARCHAR(64)", IsNullable: true, CharacterMaximumLength: &length}, want: `"code" VARCHAR(64)`}, + {name: "unsigned", column: columnInfo{Name: "count", DataType: "integer", FullDataType: "integer unsigned", IsNullable: true}, want: `"count" integer unsigned`}, + {name: "enum values", column: columnInfo{Name: "status", DataType: "enum", FullDataType: "enum('new','done')", IsNullable: true}, want: `"status" enum('new','done')`}, + {name: "datetime precision", column: columnInfo{Name: "created_at", DataType: "datetime", FullDataType: "datetime(6)", IsNullable: true}, want: `"created_at" datetime(6)`}, + {name: "bit length", column: columnInfo{Name: "mask", DataType: "bit", FullDataType: "bit(8)", IsNullable: true}, want: `"mask" bit(8)`}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if got := columnDDLDefinition(test.column); got != test.want { + t.Fatalf("unexpected column DDL: got %q, want %q", got, test.want) + } + }) + } +} + +func TestInformationSchemaColumnsPreserveFullTypesInDDL(t *testing.T) { + state := &metadataDriverState{query: func(query string) (driver.Rows, error) { + if !strings.Contains(query, "c.column_type") { + return nil, errors.New("full column type was not requested") + } + return &valueRows{ + columns: []string{"column_name", "data_type", "column_type", "is_nullable", "column_default", "column_comment", "numeric_precision", "numeric_scale", "character_maximum_length"}, + rows: [][]driver.Value{ + {"count", "integer", "integer unsigned", "YES", nil, nil, int64(32), int64(0), nil}, + {"status", "enum", "enum('new','done')", "YES", nil, nil, nil, nil, nil}, + {"created_at", "datetime", "datetime(6)", "YES", nil, nil, nil, nil, nil}, + {"mask", "bit", "bit(8)", "YES", nil, nil, nil, nil, nil}, + }, + }, nil + }} + server := newServer() + server.db = openMetadataDB(t, state) + + columns, err := server.informationSchemaColumns("public", "orders", map[string]bool{}) + if err != nil { + t.Fatal(err) + } + if len(columns) != 4 || columns[0].DataType != "integer" || columns[0].FullDataType != "integer unsigned" { + t.Fatalf("unexpected metadata columns: %#v", columns) + } + ddl := renderTableDDL("public", "orders", columns, nil) + for _, expected := range []string{ + `"count" integer unsigned`, + `"status" enum('new','done')`, + `"created_at" datetime(6)`, + `"mask" bit(8)`, + } { + if !strings.Contains(ddl, expected) { + t.Fatalf("table DDL missing %q:\n%s", expected, ddl) + } + } +} + +func TestInformationSchemaColumnsFallsBackOnlyForMissingColumnType(t *testing.T) { + tests := []struct { + name string + firstError error + wantFallback bool + }{ + {name: "missing column_type", firstError: &gokb.Error{Code: gokb.ErrorCode("42703"), Message: "column c.column_type does not exist"}, wantFallback: true}, + {name: "different missing column", firstError: &gokb.Error{Code: gokb.ErrorCode("42703"), Message: "column c.other_column does not exist"}}, + {name: "other metadata error", firstError: errors.New("metadata connection reset")}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + state := &metadataDriverState{query: func(query string) (driver.Rows, error) { + if strings.Contains(query, "c.column_type") { + return nil, test.firstError + } + if !strings.Contains(query, "NULL AS column_type") { + return nil, errors.New("unexpected fallback query: " + query) + } + return &valueRows{ + columns: []string{"column_name", "data_type", "column_type", "is_nullable", "column_default", "column_comment", "numeric_precision", "numeric_scale", "character_maximum_length"}, + rows: [][]driver.Value{{"label", "varchar", nil, "YES", nil, nil, nil, nil, int64(64)}}, + }, nil + }} + server := newServer() + server.db = openMetadataDB(t, state) + + columns, err := server.informationSchemaColumns("public", "orders", map[string]bool{}) + if test.wantFallback { + if err != nil { + t.Fatal(err) + } + if len(columns) != 1 || columnDDLDefinition(columns[0]) != `"label" varchar(64)` { + t.Fatalf("unexpected fallback columns: %#v", columns) + } + } else if !errors.Is(err, test.firstError) { + t.Fatalf("unexpected error: %v", err) + } + + state.mu.Lock() + queries := append([]string(nil), state.queries...) + state.mu.Unlock() + wantQueries := 1 + if test.wantFallback { + wantQueries = 2 + } + if len(queries) != wantQueries { + t.Fatalf("unexpected query count: got %d, want %d: %v", len(queries), wantQueries, queries) + } + }) + } +} + func TestAppendDDLStatementEnsuresSingleTerminator(t *testing.T) { got := appendDDLStatement("CREATE TABLE \"public\".\"orders\" (\n \"id\" integer\n)\n", "CREATE INDEX orders_id_idx ON public.orders (id)") want := "CREATE TABLE \"public\".\"orders\" (\n \"id\" integer\n);\n\nCREATE INDEX orders_id_idx ON public.orders (id);"