Add unit tests for clickhouse factory

Signed-off-by: haanhvu <[email protected]>

plugin/storage/clickhouse/config.go

@@ -44,17 +44,16 @@ func (cfg *Config) NewClient(ctx context.Context) (*sql.DB, error) {
  queryParams.Set("secure", "true")
  }
 
- if cfg.Database != "" {
- dsnURL.Path = cfg.Database
- } else {
- dsnURL.Path = defaultDatabase
+ if cfg.Database == "" {
+ cfg.Database = defaultDatabase
  }
+ dsnURL.Path = cfg.Database
 
- if cfg.Username != "" {
+ if cfg.Username == "" {
  cfg.Username = defaultUsername
  }
 
- if cfg.Password != "" {
+ if cfg.Password == "" {
  cfg.Password = defaultPassword
  }
 
@@ -68,7 +67,7 @@ func (cfg *Config) NewClient(ctx context.Context) (*sql.DB, error) {
  return nil, err
  }
 
- db, err := sql.Open("clickhouse", dsn)
+ db, err := sql.Open("", dsn)
  if err != nil {
  return nil, err
  }

plugin/storage/clickhouse/factory.go

@@ -40,10 +40,11 @@ func NewFactory(ctx context.Context, cfg Config, logger *zap.Logger) *Factory {
  //TODO: Move some steps to Initialize()
 }
 
-func (f *Factory) CreateSpansTable(ctx context.Context) {
+func (f *Factory) CreateSpansTable(ctx context.Context) error {
  if err := chSpanStore.CreateSpansTable(ctx, f.client, f.spansTableName); err != nil {
  f.logger.Error("failed to create spans table", zap.Error(err))
  }
+ return nil
 }
 
 func (f *Factory) ExportSpans(ctx context.Context, td ptrace.Traces) error {

plugin/storage/clickhouse/factory_test.go

@@ -0,0 +1,142 @@
+// Copyright (c) 2023 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package clickhouse
+
+import (
+ "context"
+ "database/sql"
+ "database/sql/driver"
+ "strings"
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/require"
+ "go.opentelemetry.io/collector/pdata/pcommon"
+ "go.opentelemetry.io/collector/pdata/ptrace"
+
+ "go.uber.org/zap"
+)
+
+func TestExportSpans(t *testing.T) {
+ t.Run("create factory and export spans", func(t *testing.T) {
+ var items int
+ initClickhouseTestServer(t, func(query string, values []driver.Value) error {
+ if strings.HasPrefix(query, "INSERT") {
+ items++
+ require.Equal(t, "test-operation", values[4])
+ require.Equal(t, "test-service", values[5])
+ require.Equal(t, []string{"attKey0", "attKey1"}, values[6])
+ require.Equal(t, []string{"attVal0", "attVal1"}, values[7])
+ }
+ return nil
+ })
+
+ c := Config{
+ Endpoint: "clickhouse://127.0.0.1:9000",
+ SpansTableName: "jaeger_spans",
+ }
+
+ f := NewFactory(context.TODO(), c, zap.NewNop())
+ require.NotNil(t, f.client)
+
+ err := f.CreateSpansTable(context.TODO())
+ require.Nil(t, err)
+
+ err = f.ExportSpans(context.TODO(), createTraces(5))
+ require.Nil(t, err)
+
+ err = f.ExportSpans(context.TODO(), createTraces(10))
+ require.Nil(t, err)
+
+ require.Equal(t, 15, items)
+ })
+}
+
+func createTraces(count int) ptrace.Traces {
+ traces := ptrace.NewTraces()
+ rs := traces.ResourceSpans().AppendEmpty()
+ rs.Resource().Attributes().PutStr("service.name", "test-service")
+ ss := rs.ScopeSpans().AppendEmpty()
+ for i := 0; i < count; i++ {
+ s := ss.Spans().AppendEmpty()
+ s.SetStartTimestamp(pcommon.NewTimestampFromTime(time.Now()))
+ s.SetEndTimestamp(pcommon.NewTimestampFromTime(time.Now()))
+ s.Attributes().PutStr("attKey0", "attVal0")
+ s.Attributes().PutStr("attKey1", "attVal1")
+ s.SetName("test-operation")
+ }
+ return traces
+}
+
+func initClickhouseTestServer(t *testing.T, recorder recorder) {
+ sql.Register("", &testClickhouseDriver{
+ recorder: recorder,
+ })
+}
+
+type recorder func(query string, values []driver.Value) error
+
+type testClickhouseDriver struct {
+ recorder recorder
+}
+
+func (t *testClickhouseDriver) Open(_ string) (driver.Conn, error) {
+ return &testClickhouseDriverConn{
+ recorder: t.recorder,
+ }, nil
+}
+
+type testClickhouseDriverConn struct {
+ recorder recorder
+}
+
+func (*testClickhouseDriverConn) Begin() (driver.Tx, error) {
+ return &testClickhouseDriverTx{}, nil
+}
+
+func (*testClickhouseDriverConn) Close() error {
+ return nil
+}
+
+func (t *testClickhouseDriverConn) Prepare(query string) (driver.Stmt, error) {
+ return &testClickhouseDriverStmt{
+ query: query,
+ recorder: t.recorder,
+ }, nil
+}
+
+func (*testClickhouseDriverConn) CheckNamedValue(_ *driver.NamedValue) error {
+ return nil
+}
+
+type testClickhouseDriverTx struct{}
+
+func (*testClickhouseDriverTx) Commit() error {
+ return nil
+}
+
+func (*testClickhouseDriverTx) Rollback() error {
+ return nil
+}
+
+type testClickhouseDriverStmt struct {
+ query string
+ recorder recorder
+}
+
+func (*testClickhouseDriverStmt) Close() error {
+ return nil
+}
+
+func (t *testClickhouseDriverStmt) Exec(args []driver.Value) (driver.Result, error) {
+ return nil, t.recorder(t.query, args)
+}
+
+func (t *testClickhouseDriverStmt) NumInput() int {
+ return strings.Count(t.query, "?")
+}
+
+func (t *testClickhouseDriverStmt) Query(_ []driver.Value) (driver.Rows, error) {
+ return nil, nil
+}

plugin/storage/clickhouse/spanstore/exporter.go

@@ -15,8 +15,7 @@ import (
 )
 
 const (
- insertSpansSQL = `
- INSERT INTO %s (
+ insertSpansSQL = `INSERT INTO %s (
  Timestamp,
  TraceId,
  SpanId,
@@ -35,7 +34,7 @@ const (
  ?,
  ?,
  ?,
- ?
+ ?,
  )
  `
 )
@@ -91,7 +90,7 @@ func insertSpans(ctx context.Context, tx *sql.Tx, tableName string, td ptrace.Tr
  uint64(r.EndTimestamp().AsTime().Sub(r.StartTimestamp().AsTime()).Nanoseconds()),
  )
  if err != nil {
- return err
+ return fmt.Errorf("exec context: %s", err)
  }
  }
  }

plugin/storage/clickhouse/spanstore/schema.go

@@ -10,8 +10,7 @@ import (
 )
 
 const (
- createSpansTableSQL = `
- CREATE TABLE IF NOT EXISTS %s (
+ createSpansTableSQL = `CREATE TABLE IF NOT EXISTS %s (
  Timestamp DateTime64(9) CODEC(Delta, ZSTD(1)),
  TraceId String CODEC(ZSTD(1)),
  SpanId String CODEC(ZSTD(1)),