cmd/apply_comments.go

/*
 * Copyright 2025 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package cmd

import (
	"fmt"
	"log"

	"github.com/GoogleCloudPlatform/db-context-enrichment/internal/database"
	"github.com/GoogleCloudPlatform/db-context-enrichment/internal/utils"
	"github.com/spf13/cobra"
)

var applyCommentsCmd = &cobra.Command{
	Use:     "apply-comments",
	Short:   "Apply comments to database columns from a SQL file",
	Long:    `Reads SQL statements from a specified file and applies them to the database. This is typically used to apply comments generated by the 'add-comments' command.`,
	Example: `db_schema_enricher apply-comments --dialect cloudsqlpostgres --username user --password pass --database mydb --cloudsql-instance-connection-name my-project:my-region:my-instance --in_file ./column_comments.sql`,
	RunE: func(cmd *cobra.Command, args []string) error {
		db, err := setupDatabase()
		if err != nil {
			return err
		}
		defer db.Close()
		inputFile := cmd.Flag("in_file").Value.String() // Get inputFile here
		return runApplyComments(cmd, inputFile, db)
	},
}

func runApplyComments(cmd *cobra.Command, sqlFilePath string, db *database.DB) error {
	dbConfig := database.GetConfig()
	if dbConfig == nil {
		return fmt.Errorf("database config is not initialized")
	}
	//Set default input file if not provided
	inputFile := sqlFilePath // Use the function parameter
	if inputFile == "" {
		inputFile = fmt.Sprintf("%s_comments.sql", dbConfig.DBName) // Default input file name
	}
	log.Println("INFO: Starting apply-comments operation from file:", inputFile)

	sqlStatements, err := utils.ReadSQLStatementsFromFile(inputFile)
	if err != nil {
		return fmt.Errorf("failed to read SQL statements from file: %w", err)
	}

	ctx := cmd.Context()
	if execErr := db.ExecuteSQLStatements(ctx, sqlStatements); execErr != nil {
		return fmt.Errorf("failed to apply comments from SQL file: %w", execErr)
	}

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

func init() {
	var inputFile string

	// Flags for apply-comments command
	applyCommentsCmd.Flags().StringVarP(&inputFile, "in_file", "i", "", "File path to read SQL statements from (defaults to <database>_comments.sql)")
}