func()

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
}