From 0cb78c5fd49c35a76b58afb82aebb37a1e047f09 Mon Sep 17 00:00:00 2001
From: Boris Tyshkevich <68195949+bvt123@users.noreply.github.com>
Date: Tue, 10 Jun 2025 22:22:43 +0200
Subject: [PATCH] Remove clickhouse infer query workaround

---
 connectors/base.go            |  6 ++++++
 connectors/init.go            |  6 +++++-
 connectors/logging_wrapper.go | 36 +++++++++++++++++++++++++++++++++++
 main.go                       | 13 +++++++++++++
 4 files changed, 60 insertions(+), 1 deletion(-)
 create mode 100644 connectors/logging_wrapper.go

diff --git a/connectors/base.go b/connectors/base.go
index 4da20f6..287992e 100644
--- a/connectors/base.go
+++ b/connectors/base.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/centralmind/gateway/model"
 	"github.com/jmoiron/sqlx"
+	"github.com/sirupsen/logrus"
 	"golang.org/x/xerrors"
 )
 
@@ -27,15 +28,18 @@ func (b *BaseConnector) mapSQLTypeToColumnType(guesser TypeGuesser, sqlType stri
 // InferResultColumns provides a generic implementation for getting result column information
 // This implementation works with any SQL database that supports the database/sql interfaces
 func (b *BaseConnector) InferResultColumns(ctx context.Context, query string, guesser TypeGuesser) ([]model.ColumnSchema, error) {
+	logrus.Infof("Infer columns for query: %s", query)
 	// Prepare the statement to get column information
 	tx, err := b.DB.BeginTxx(ctx, &sql.TxOptions{
 		ReadOnly: true,
 	})
 	if err != nil {
+		logrus.Errorf("BeginTx failed with error: %v", err)
 		return nil, xerrors.Errorf("BeginTx failed with error: %w", err)
 	}
 	stmt, err := tx.PrepareNamedContext(ctx, query)
 	if err != nil {
+		logrus.Errorf("unable to prepare statement: %v", err)
 		return nil, xerrors.Errorf("unable to prepare statement: %w", err)
 	}
 	defer stmt.Close()
@@ -49,6 +53,7 @@ func (b *BaseConnector) InferResultColumns(ctx context.Context, query string, gu
 	// to only fetch metadata without actually executing the full query
 	rows, err := stmt.QueryContext(ctx, prms)
 	if err != nil {
+		logrus.Errorf("unable to execute statement: %v", err)
 		return nil, xerrors.Errorf("unable to execute statement: %w", err)
 	}
 	defer rows.Close()
@@ -56,6 +61,7 @@ func (b *BaseConnector) InferResultColumns(ctx context.Context, query string, gu
 	// Get column types from the result
 	colTypes, err := rows.ColumnTypes()
 	if err != nil {
+		logrus.Errorf("unable to get column types: %v", err)
 		return nil, xerrors.Errorf("unable to get column types: %w", err)
 	}
 
diff --git a/connectors/init.go b/connectors/init.go
index 6187390..6d95c95 100644
--- a/connectors/init.go
+++ b/connectors/init.go
@@ -76,5 +76,9 @@ func New(tag string, config any) (Connector, error) {
 	if !ok {
 		return nil, xerrors.Errorf("connector: %s not found", tag)
 	}
-	return f(config)
+	conn, err := f(config)
+	if err != nil {
+		return nil, err
+	}
+	return WrapWithLogging(conn), nil
 }
diff --git a/connectors/logging_wrapper.go b/connectors/logging_wrapper.go
new file mode 100644
index 0000000..683b4c2
--- /dev/null
+++ b/connectors/logging_wrapper.go
@@ -0,0 +1,36 @@
+package connectors
+
+import (
+	"context"
+
+	"github.com/sirupsen/logrus"
+
+	"github.com/centralmind/gateway/model"
+)
+
+type loggingConnector struct {
+	Connector
+}
+
+// WrapWithLogging wraps the given connector with logging functionality.
+func WrapWithLogging(inner Connector) Connector {
+	return &loggingConnector{Connector: inner}
+}
+
+func (l loggingConnector) Query(ctx context.Context, endpoint model.Endpoint, params map[string]any) ([]map[string]any, error) {
+	logrus.Infof("SQL query: %s", endpoint.Query)
+	res, err := l.Connector.Query(ctx, endpoint, params)
+	if err != nil {
+		logrus.Errorf("query failed: %v", err)
+	}
+	return res, err
+}
+
+func (l loggingConnector) InferQuery(ctx context.Context, query string) ([]model.ColumnSchema, error) {
+	logrus.Infof("Infer query: %s", query)
+	res, err := l.Connector.InferQuery(ctx, query)
+	if err != nil {
+		logrus.Errorf("infer query failed: %v", err)
+	}
+	return res, err
+}
diff --git a/main.go b/main.go
index 9d1e418..61eff2c 100644
--- a/main.go
+++ b/main.go
@@ -1,8 +1,10 @@
 package main
 
 import (
+	"fmt"
 	"os"
 
+	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 
 	"github.com/centralmind/gateway/cli"
@@ -34,6 +36,17 @@ func main() {
 		Example:      "./gateway help",
 		SilenceUsage: true,
 	}
+	var logLevel string
+	rootCommand.PersistentFlags().StringVar(&logLevel, "log-level", "info", "logging level (trace, debug, info, warn, error)")
+	rootCommand.PersistentPreRunE = func(cmd *cobra.Command, args []string) error {
+		level, err := logrus.ParseLevel(logLevel)
+		if err != nil {
+			return fmt.Errorf("invalid log level %s", logLevel)
+		}
+		logrus.SetOutput(os.Stdout)
+		logrus.SetLevel(level)
+		return nil
+	}
 	cli.RegisterCommand(rootCommand, cli.StartCommand())
 	cli.RegisterCommand(rootCommand, cli.Connectors())
 	cli.RegisterCommand(rootCommand, cli.Plugins())