in internal/enricher/enricher.go [106:184]
func (mc *MetadataCollector) CollectColumnMetadata(ctx context.Context, tableName string, colInfo database.ColumnInfo, schemaContext string) (*ColumnMetadata, error) {
if tableName == "" || colInfo.Name == "" {
return nil, &ErrInvalidInput{
Msg: "table name and column name cannot be empty",
}
}
var (
exampleValues []string
distinctCount int64
nullCount int64
err error
)
dbMetadata := make(map[string]interface{})
if mc.isEnrichmentRequested("examples") || mc.isEnrichmentRequested("distinct_values") || mc.isEnrichmentRequested("null_count") {
dbMetadata, err = mc.db.GetColumnMetadata(tableName, colInfo.Name)
if err != nil {
return nil, &ErrQueryExecution{
Msg: "failed to get column metadata",
Err: err,
}
}
}
if mc.isEnrichmentRequested("examples") {
retrivedExampleValues, ok := dbMetadata["ExampleValues"].([]string)
if !ok {
log.Println("WARN: unexpected type for ExampleValues")
}
exampleValues = retrivedExampleValues
if mc.GeminiAPIKey != "" {
processedExampleValues, err := mc.generateExampleValuesWithGemini(ctx, colInfo, tableName, exampleValues)
if err != nil {
log.Printf("WARN: Failed to generate/process example values with Gemini for column %s.%s: %v, using original examples.", tableName, colInfo.Name, err)
} else {
exampleValues = processedExampleValues
}
}
}
if mc.isEnrichmentRequested("distinct_values") {
distinctCountFloat, ok := dbMetadata["DistinctCount"].(int)
if !ok {
log.Println("WARN: unexpected type for DistinctCount")
}
distinctCount = int64(distinctCountFloat)
}
if mc.isEnrichmentRequested("null_count") {
nullCountFloat, ok := dbMetadata["NullCount"].(int)
if !ok {
log.Println("WARN: unexpected type for NullCount")
}
nullCount = int64(nullCountFloat)
}
metadata := &ColumnMetadata{
Table: tableName,
Column: colInfo.Name,
DataType: colInfo.DataType,
ExampleValues: exampleValues,
DistinctCount: distinctCount,
NullCount: nullCount,
}
// Generate description using Gemini if API key is available and "description" enrichment is requested
if mc.GeminiAPIKey != "" && mc.isEnrichmentRequested("description") {
description, err := mc.generateDescriptionWithGemini(ctx, metadata, schemaContext)
if err != nil {
log.Printf("WARN: Failed to generate description with Gemini for %s.%s: %v", tableName, colInfo.Name, err)
} else {
metadata.Description = description
}
}
return metadata, nil
}