github.com/erda-project/erda-infra@v1.0.9/providers/clickhouse/clickhouse.go (about)

     1  // Copyright (c) 2021 Terminus, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package clickhouse
    16  
    17  import (
    18  	"fmt"
    19  	"reflect"
    20  	"strings"
    21  	"time"
    22  
    23  	ck "github.com/ClickHouse/clickhouse-go/v2"
    24  	ckdriver "github.com/ClickHouse/clickhouse-go/v2/lib/driver"
    25  
    26  	"github.com/erda-project/erda-infra/base/logs"
    27  	"github.com/erda-project/erda-infra/base/servicehub"
    28  )
    29  
    30  // Interface clickhouse client
    31  type Interface interface {
    32  
    33  	// Client .
    34  	Client() ckdriver.Conn
    35  
    36  	// NewWriter .
    37  	NewWriter(opts *WriterOptions) *Writer
    38  }
    39  
    40  var (
    41  	interfaceType  = reflect.TypeOf((*Interface)(nil)).Elem()
    42  	nativeConnType = reflect.TypeOf((*ckdriver.Conn)(nil)).Elem()
    43  )
    44  
    45  type config struct {
    46  	Addr             string        `file:"addr" default:"localhost:9000"`
    47  	Username         string        `file:"username" default:"default"`
    48  	Password         string        `file:"password"`
    49  	Database         string        `file:"database"`
    50  	DialTimeout      time.Duration `file:"dial_timeout" default:"1s"`
    51  	MaxIdleConns     int           `file:"max_idle_conns" default:"5"`
    52  	MaxOpenConns     int           `file:"max_open_conns" default:"10"`
    53  	ConnMaxLifeTime  time.Duration `file:"conn_max_lifetime" default:"1h"`
    54  	ConnOpenStrategy string        `file:"conn_open_strategy" default:"in_order"`
    55  	Debug            bool          `file:"debug"`
    56  }
    57  
    58  type provider struct {
    59  	Cfg *config
    60  	Log logs.Logger
    61  
    62  	nativeConn ckdriver.Conn
    63  }
    64  
    65  func (p *provider) Init(ctx servicehub.Context) error {
    66  	openStrategy := ck.ConnOpenInOrder
    67  	switch p.Cfg.ConnOpenStrategy {
    68  	case "0", "in_order":
    69  		openStrategy = ck.ConnOpenInOrder
    70  	case "1", "round_robin":
    71  		openStrategy = ck.ConnOpenRoundRobin
    72  	}
    73  
    74  	options := &ck.Options{
    75  		Addr: strings.Split(p.Cfg.Addr, ","),
    76  		Auth: ck.Auth{
    77  			Database: p.Cfg.Database,
    78  			Username: p.Cfg.Username,
    79  			Password: p.Cfg.Password,
    80  		},
    81  		DialTimeout:      p.Cfg.DialTimeout,
    82  		Debug:            p.Cfg.Debug,
    83  		MaxIdleConns:     p.Cfg.MaxIdleConns,
    84  		MaxOpenConns:     p.Cfg.MaxOpenConns,
    85  		ConnMaxLifetime:  p.Cfg.ConnMaxLifeTime,
    86  		ConnOpenStrategy: openStrategy,
    87  	}
    88  
    89  	conn, err := ck.Open(options)
    90  	if err != nil {
    91  		return fmt.Errorf("fail to connect clickhouse: %s", err)
    92  	}
    93  	p.nativeConn = conn
    94  
    95  	return nil
    96  }
    97  
    98  func (p *provider) Client() ckdriver.Conn {
    99  	return p.nativeConn
   100  }
   101  
   102  func (p *provider) NewWriter(opts *WriterOptions) *Writer {
   103  	return NewWriter(p.nativeConn, opts.Encoder)
   104  }
   105  
   106  func (p *provider) Provide(ctx servicehub.DependencyContext, args ...interface{}) interface{} {
   107  	if ctx.Service() == "clickhouse-client" || ctx.Type() == nativeConnType {
   108  		return p.nativeConn
   109  	}
   110  	return p
   111  }
   112  
   113  func init() {
   114  	servicehub.Register("clickhouse", &servicehub.Spec{
   115  		Services:    []string{"clickhouse", "clickhouse-client"},
   116  		Types:       []reflect.Type{interfaceType, nativeConnType},
   117  		Description: "clickhouse client",
   118  		ConfigFunc: func() interface{} {
   119  			return &config{}
   120  		},
   121  		Creator: func() servicehub.Provider {
   122  			return &provider{}
   123  		},
   124  	})
   125  }