fix(kingbase): handle missing type metadata columns

Closes #5173

agents/drivers/kingbase-go/kingbase_metadata.go

@@ -539,24 +539,42 @@ func isUndefinedColumn(err error, columnName string) bool {
 }
 
 func (s *server) informationSchemaColumns(schema, table string, primary map[string]bool) ([]columnInfo, error) {
-	result, err := s.queryInformationSchemaColumns(schema, table, primary, true)
-	if err != nil && isUndefinedColumn(err, "column_type") {
-		return s.queryInformationSchemaColumns(schema, table, primary, false)
+	includeColumnType := true
+	includeUdtName := true
+	for {
+		result, err := s.queryInformationSchemaColumns(schema, table, primary, includeColumnType, includeUdtName)
+		if err == nil {
+			return result, nil
+		}
+		switch {
+		case includeColumnType && isUndefinedColumn(err, "column_type"):
+			includeColumnType = false
+		case includeUdtName && isUndefinedColumn(err, "udt_name"):
+			includeUdtName = false
+		default:
+			return nil, err
+		}
 	}
-	return result, err
 }
 
-func (s *server) queryInformationSchemaColumns(schema, table string, primary map[string]bool, includeFullDataType bool) ([]columnInfo, error) {
-	fullDataTypeExpression := `CASE
+func (s *server) queryInformationSchemaColumns(schema, table string, primary map[string]bool, includeColumnType, includeUdtName bool) ([]columnInfo, error) {
+	var fullDataTypeExpression string
+	switch {
+	case includeColumnType && includeUdtName:
+		fullDataTypeExpression = `CASE
 	WHEN UPPER(TRIM(c.data_type)) IN ('USER-DEFINED', 'USER_DEFINED')
 		AND UPPER(COALESCE(NULLIF(TRIM(c.column_type), ''), 'USER-DEFINED')) IN ('USER-DEFINED', 'USER_DEFINED')
 	THEN c.udt_name
 	ELSE c.column_type
 	END`
-	if !includeFullDataType {
+	case includeColumnType:
+		fullDataTypeExpression = "c.column_type"
+	case includeUdtName:
 		fullDataTypeExpression = `CASE
 		WHEN UPPER(TRIM(c.data_type)) IN ('USER-DEFINED', 'USER_DEFINED') THEN c.udt_name
 		END AS column_type`
+	default:
+		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

agents/drivers/kingbase-go/main_test.go

@@ -1158,7 +1158,73 @@ func TestInformationSchemaColumnsResolveUserDefinedTypeWithoutColumnType(t *test
 	}
 }
 
-func TestInformationSchemaColumnsFallsBackOnlyForMissingColumnType(t *testing.T) {
+func TestInformationSchemaColumnsPreserveColumnTypeWithoutUdtName(t *testing.T) {
+	state := &metadataDriverState{query: func(query string) (driver.Rows, error) {
+		if strings.Contains(query, "c.udt_name") {
+			return nil, &gokb.Error{Code: gokb.ErrorCode("42703"), Message: "kb: column c.udt_name does not exist"}
+		}
+		if !strings.Contains(query, "c.column_type") {
+			return nil, errors.New("column type fallback 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{{"status", "enum", "enum('new','done')", "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) != 1 || columns[0].FullDataType != "enum('new','done')" {
+		t.Fatalf("unexpected metadata columns: %#v", columns)
+	}
+	if ddl := renderTableDDL("public", "orders", columns, nil); !strings.Contains(ddl, `"status" enum('new','done')`) {
+		t.Fatalf("unexpected table DDL:\n%s", ddl)
+	}
+	state.mu.Lock()
+	defer state.mu.Unlock()
+	if len(state.queries) != 2 {
+		t.Fatalf("unexpected query count: got %d, want 2: %v", len(state.queries), state.queries)
+	}
+}
+
+func TestInformationSchemaColumnsFallbackWithoutExtendedTypeColumns(t *testing.T) {
+	state := &metadataDriverState{query: func(query string) (driver.Rows, error) {
+		if strings.Contains(query, "c.column_type") {
+			return nil, &gokb.Error{Code: gokb.ErrorCode("42703"), Message: "column c.column_type does not exist"}
+		}
+		if strings.Contains(query, "c.udt_name") {
+			return nil, &gokb.Error{Code: gokb.ErrorCode("42703"), Message: "column c.udt_name does not exist"}
+		}
+		if !strings.Contains(query, "NULL AS column_type") {
+			return nil, errors.New("base type fallback 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{{"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 err != nil {
+		t.Fatal(err)
+	}
+	if len(columns) != 1 || columnDDLDefinition(columns[0]) != `"label" varchar(64)` {
+		t.Fatalf("unexpected fallback columns: %#v", columns)
+	}
+	state.mu.Lock()
+	defer state.mu.Unlock()
+	if len(state.queries) != 3 {
+		t.Fatalf("unexpected query count: got %d, want 3: %v", len(state.queries), state.queries)
+	}
+}
+
+func TestInformationSchemaColumnsRetriesOnlyForMissingTypeMetadataColumns(t *testing.T) {
 	tests := []struct {
 		name         string
 		firstError   error