func runAddComments()

in cmd/add_comments.go [40:161]
99 lines of code
25 McCabe index (conditional complexity)

func runAddComments(cmd *cobra.Command, args []string) error {
	dbConfig := database.GetConfig()
	if dbConfig == nil {
		return fmt.Errorf("database config is not initialized")
	}

	if err := validateDialect(dialect); err != nil {
		return err
	}

	outputFile := cmd.Flag("out_file").Value.String()
	if outputFile == "" {
		outputFile = utils.GetDefaultOutputFilePath(dbConfig.DBName, "add-comments")
	}

	log.Println("INFO: Starting add-comments operation", "dialect:", dbConfig.Dialect, "database:", dbConfig.DBName)

	db, err := setupDatabase()
	if err != nil {
		return err
	}
	defer db.Close()

	model := cmd.Flag("model").Value.String()
	metadataCollector := enricher.NewMetadataCollector(db, &enricher.DefaultRetryOptions, dryRun, geminiAPIKey, "", model)

	ctx := cmd.Context()

	tablesFlag := cmd.Flag("tables").Value.String()
	tableFilters, err := utils.ParseTablesFlag(tablesFlag)
	if err != nil {
		return err
	}
	metadataCollector.TableFilters = tableFilters

	enrichmentsFlag := cmd.Flag("enrichments").Value.String()
	enrichmentSet := make(map[string]bool)
	if enrichmentsFlag != "" {
		enrichmentsFlag = strings.ReplaceAll(enrichmentsFlag, " ", "")
		for _, e := range strings.Split(enrichmentsFlag, ",") {
			enrichmentSet[strings.TrimSpace(strings.ToLower(e))] = true
		}
	}
	metadataCollector.Enrichments = enrichmentSet

	contextFilesFlag := cmd.Flag("context").Value.String()
	additionalContext, err := utils.ReadContextFiles(contextFilesFlag)
	if err != nil {
		return fmt.Errorf("failed to read context files: %w", err)
	}
	metadataCollector.AdditionalContext = additionalContext

	// API Key Validation Logic
	if additionalContext != "" {
		if geminiAPIKey == "" {
			return fmt.Errorf("additional context is provided, but Gemini API key is not configured. Please set the GEMINI_API_KEY environment variable")
		}
		if err := metadataCollector.IsGeminiAPIKeyValid(ctx); err != nil {
			return fmt.Errorf("\nGemini API key is invalid. Please provide a valid api key\n")
		}
	}

	if len(enrichmentSet) == 0 || enrichmentSet["examples"] {
		if geminiAPIKey == "" {
			log.Println("WARN: No Gemini API key provided. PII identification will be skipped.")
		}
		if err := metadataCollector.IsGeminiAPIKeyValid(ctx); err != nil {
			geminiAPIKey = ""
			log.Println("WARN: Gemini API key provided is invalid. PII identification will be skipped.\n")
		}
	}

	metadataCollector.GeminiAPIKey = geminiAPIKey

	sqlStatements, err := metadataCollector.GenerateCommentSQLs(ctx)
	if err != nil {
		return fmt.Errorf("metadata collection and SQL generation failed: %w", err)
	}

	file, createErr := os.Create(outputFile)
	if createErr != nil {
		return fmt.Errorf("failed to create output file: %w", createErr)
	}
	defer file.Close()

	for _, sqlStmt := range sqlStatements {
		if _, writeErr := file.WriteString(sqlStmt + "\n"); writeErr != nil {
			return fmt.Errorf("failed to write SQL statement to file: %w", writeErr)
		}
	}

	log.Println("INFO: SQL statements to add column comments have been written to:", outputFile)

	if dryRun {
		log.Println("INFO: Add comments operation completed in dry-run mode.  No changes were made to the database.")
		return nil
	}

	// --- User Confirmation ---
	if len(sqlStatements) > 0 {
		if utils.ConfirmAction("SQL statements to add column comments") {
			// Re-read the SQL statements from the output file as changes may have been
			fileContent, readErr := os.ReadFile(outputFile)
			if readErr != nil {
				return fmt.Errorf("failed to read SQL statements from output file: %w", readErr)
			}
			sqlStatements = strings.Split(strings.TrimSpace(string(fileContent)), "\n")

			if execErr := db.ExecuteSQLStatements(ctx, sqlStatements); execErr != nil {
				return fmt.Errorf("failed to execute SQL statements to add comments: %w", execErr)
			}
			log.Println("INFO: Successfully added comments to the database.")
		} else {
			log.Println("INFO: Comment addition aborted by user.")
		}
	} else {
		log.Println("INFO: No comments to add.")
	}

	log.Println("INFO: Add comments operation completed.")
	return nil
}