basic datasets import for clickhouse

Author: sholokhov

cmd/check.go

@@ -25,7 +25,6 @@ to quickly create a Cobra application.`,
 			if dataSource != "" {
 				fmt.Println("Data source is not empty: " + dataSource)
 			}
-			app.GetDbqConfig()
 			return nil
 		},
 	}

cmd/import.go

@@ -2,7 +2,7 @@ package cmd
 
 import (
 	"dbq/internal"
-	"fmt"
+	"log"
 
 	"github.com/spf13/cobra"
 )
@@ -24,16 +24,21 @@ to quickly create a Cobra application.`,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			datasets, err := app.ImportDatasets(dataSource, filter)
 			if err != nil {
-				fmt.Println("Failed to fetch datasets: " + err.Error())
+				log.Println("Failed to fetch datasets: " + err.Error())
 				return nil
 			}
 
-			for _, dataset := range datasets {
-				fmt.Println("Imported dataset: ", dataset)
-			}
-
+			log.Printf("Found %d datasets to import: %v\n", len(datasets), datasets)
 			if updateCfg {
-				fmt.Println("Updating dbq config...")
+				ds := app.FindDataSourceById(dataSource)
+				if ds != nil {
+					ds.Datasets = datasets
+					err := app.SaveDbqConfig()
+					if err != nil {
+						return err
+					}
+					log.Println("dbq config has been updated")
+				}
 			}
 
 			return nil

cmd/ping.go

@@ -29,7 +29,6 @@ to quickly create a Cobra application.`,
 	}
 
 	cmd.Flags().StringVarP(&dataSource, "datasource", "d", "", "Datasource")
-	//_ = cmd.MarkFlagRequired("dataSource")
 
 	return cmd
 }

dbq.yaml

@@ -1,24 +1,23 @@
-version: 1
+version: "1"
 datasources:
-  - id: clickhouse-prod
-    type: clickhouse
-    configuration:
-      host: 0.0.0.0:19000
-      port: 19000
-      username: default
-      password: changeme
-      database: default
-    datasets:
-      - default.my_table_1
-      - default.my_table_2
-      - ads.my_table_1
-      - kpi.some_table
-  - id: pgsql-staging
-    type: postgres
-    configuration:
-      host: 127.0.0.1
-      port: 9004
-      username: simple
-      password: simple_pass
-    datasets:
-      - public.table_1
+    - id: ch-local
+      type: clickhouse
+      configuration:
+        host: 0.0.0.0:19000
+        port: 19000
+        username: default
+        password: changeme
+        database: default
+      datasets:
+        - default.my_first_table
+        - default.table_2
+        - default.trips_small
+    - id: pgsql
+      type: postgres
+      configuration:
+        host: 127.0.0.1
+        port: 9004
+        username: simple
+        password: simple_pass
+      datasets:
+        - public.table_1

internal/app.go

@@ -3,32 +3,65 @@ package internal
 import (
 	"errors"
 	"fmt"
+	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
+	"gopkg.in/yaml.v3"
+	"log"
+	"os"
 	"strings"
 )
 
 type DbqApp interface {
 	PingDataSource(srcId string) error
 	ImportDatasets(srcId string, filter string) ([]string, error)
 	GetDbqConfig() *DbqConfig
+	SaveDbqConfig() error
+	FindDataSourceById(srcId string) *DataSource
 }
 
 type DbqAppImpl struct {
-	dbqConfig *DbqConfig
+	dbqConfigPath string
+	dbqConfig     *DbqConfig
 }
 
-func NewDbqApp(dbqConfig *DbqConfig) DbqApp {
-	return &DbqAppImpl{dbqConfig: dbqConfig}
+func NewDbqApp(dbqConfigPath string) DbqApp {
+	v := viper.New()
+
+	if dbqConfigPath != "" {
+		v.SetConfigFile(dbqConfigPath)
+	} else {
+		home, err := os.UserHomeDir()
+		cobra.CheckErr(err)
+		v.AddConfigPath(home)
+		v.SetConfigType("yaml")
+		v.SetConfigName(".dbq.yaml")
+	}
+
+	v.AutomaticEnv()
+	if err := v.ReadInConfig(); err != nil {
+		cobra.CheckErr(err)
+	}
+
+	var dbqConfig DbqConfig
+	if err := v.Unmarshal(&dbqConfig); err != nil {
+		cobra.CheckErr(err)
+	}
+
+	return &DbqAppImpl{
+		dbqConfigPath: dbqConfigPath,
+		dbqConfig:     &dbqConfig,
+	}
 }
 
 func (app *DbqAppImpl) PingDataSource(srcId string) error {
-	var dataSource = findDataSourceById(srcId, app.dbqConfig.DataSources)
+	var dataSource = app.FindDataSourceById(srcId)
 
 	cnn, err := getDbqConnector(*dataSource)
 	if err != nil {
 		return err
 	}
 
-	fmt.Println("Pinging datasource: " + dataSource.ID)
+	log.Println("Pinging datasource: " + dataSource.ID)
 	err = cnn.Ping()
 	if err != nil {
 		return err
@@ -38,7 +71,7 @@ func (app *DbqAppImpl) PingDataSource(srcId string) error {
 }
 
 func (app *DbqAppImpl) ImportDatasets(srcId string, filter string) ([]string, error) {
-	var dataSource = findDataSourceById(srcId, app.dbqConfig.DataSources)
+	var dataSource = app.FindDataSourceById(srcId)
 	cnn, err := getDbqConnector(*dataSource)
 	if err != nil {
 		return []string{}, err
@@ -51,6 +84,29 @@ func (app *DbqAppImpl) GetDbqConfig() *DbqConfig {
 	return app.dbqConfig
 }
 
+func (app *DbqAppImpl) SaveDbqConfig() error {
+	updatedYaml, err := yaml.Marshal(app.dbqConfig)
+	if err != nil {
+		return err
+	}
+
+	err = os.WriteFile(app.dbqConfigPath, updatedYaml, 0644)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (app *DbqAppImpl) FindDataSourceById(srcId string) *DataSource {
+	for i := range app.dbqConfig.DataSources {
+		if app.dbqConfig.DataSources[i].ID == srcId {
+			return &app.dbqConfig.DataSources[i]
+		}
+	}
+	return nil
+}
+
 func getDbqConnector(ds DataSource) (DbqConnector, error) {
 	dsType := strings.ToLower(ds.Type)
 	switch dsType {
@@ -60,12 +116,3 @@ func getDbqConnector(ds DataSource) (DbqConnector, error) {
 		return nil, errors.New(fmt.Sprintf("Data source type '%s' is not supported.", dsType))
 	}
 }
-
-func findDataSourceById(srcId string, dataSources []DataSource) *DataSource {
-	for _, src := range dataSources {
-		if src.ID == srcId {
-			return &src
-		}
-	}
-	return nil
-}

internal/cfg.go internal/checks_config.gointernal/cfg.go renamed to internal/checks_config.go

@@ -1,8 +1,8 @@
 package internal
 
 import (
-	"fmt"
 	"gopkg.in/yaml.v3"
+	"log"
 	"os"
 )
 
@@ -36,14 +36,14 @@ func LoadChecksConfig(fileName string) (*ChecksConfig, error) {
 	file, err := os.Open(fileName)
 	defer file.Close()
 	if err != nil {
-		fmt.Printf("Error opening file: %v\n", err)
+		log.Printf("Error opening file: %v\n", err)
 		return nil, err
 	}
 
 	var cfg ChecksConfig
 	decoder := yaml.NewDecoder(file)
 	if err := decoder.Decode(&cfg); err != nil {
-		fmt.Printf("Error decoding YAML: %v\n", err)
+		log.Printf("Error decoding YAML: %v\n", err)
 		return nil, err
 	}
 

internal/clickhouse.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"github.com/ClickHouse/clickhouse-go/v2"
 	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
-	"log"
 )
 
 type ClickhouseDbqConnector struct {
@@ -63,7 +62,6 @@ func (c *ClickhouseDbqConnector) ImportDatasets() ([]string, error) {
 		return nil, fmt.Errorf("error occurred during row iteration: %w", err)
 	}
 
-	log.Printf("Found %d datasets (tables).\n", len(datasets))
 	return datasets, nil
 }
 

internal/dbq_config.go

@@ -1,8 +1,8 @@
 package internal
 
 import (
-	"fmt"
 	"gopkg.in/yaml.v3"
+	"log"
 	"os"
 )
 
@@ -22,14 +22,14 @@ func LoadDbqSetting(fileName string) (*DbqConfig, error) {
 	file, err := os.Open(fileName)
 	defer file.Close()
 	if err != nil {
-		fmt.Printf("Error opening file: %v\n", err)
+		log.Printf("Error opening file: %v\n", err)
 		return nil, err
 	}
 
 	var settings DbqConfig
 	decoder := yaml.NewDecoder(file)
 	if err := decoder.Decode(&settings); err != nil {
-		fmt.Printf("Error decoding YAML: %v\n", err)
+		log.Printf("Error decoding YAML: %v\n", err)
 		return nil, err
 	}
 

main.go

@@ -8,7 +8,6 @@ import (
 	"dbq/internal"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
-	"github.com/spf13/viper"
 	"os"
 )
 
@@ -21,35 +20,8 @@ func main() {
 		cobra.CheckErr(err)
 	}
 
-	dbqConfig := initSettings(*dbqConfigFile)
-	app := internal.NewDbqApp(dbqConfig)
+	app := internal.NewDbqApp(*dbqConfigFile)
 
 	cmd.AddCommands(app)
 	cmd.Execute()
 }
-
-func initSettings(dbqConfigPath string) *internal.DbqConfig {
-	v := viper.New()
-
-	if dbqConfigPath != "" {
-		v.SetConfigFile(dbqConfigPath)
-	} else {
-		home, err := os.UserHomeDir()
-		cobra.CheckErr(err)
-		v.AddConfigPath(home)
-		v.SetConfigType("yaml")
-		v.SetConfigName(".dbq.yaml")
-	}
-
-	v.AutomaticEnv()
-	if err := v.ReadInConfig(); err != nil {
-		cobra.CheckErr(err)
-	}
-
-	var settings internal.DbqConfig
-	if err := v.Unmarshal(&settings); err != nil {
-		cobra.CheckErr(err)
-	}
-
-	return &settings
-}