/vendor/github.com/ClickHouse/clickhouse-go/clickhouse.go

https://github.com/rudderlabs/rudder-server · Go · 335 lines · 309 code · 23 blank · 3 comment · 60 complexity · 0f5e7d6fafec5ecb0896857e82a7ce63 MD5 · raw file

  1. package clickhouse
  2. import (
  3. "bufio"
  4. "context"
  5. "database/sql"
  6. "database/sql/driver"
  7. "errors"
  8. "fmt"
  9. "net"
  10. "reflect"
  11. "regexp"
  12. "sync"
  13. "time"
  14. "github.com/ClickHouse/clickhouse-go/lib/binary"
  15. "github.com/ClickHouse/clickhouse-go/lib/column"
  16. "github.com/ClickHouse/clickhouse-go/lib/data"
  17. "github.com/ClickHouse/clickhouse-go/lib/protocol"
  18. "github.com/ClickHouse/clickhouse-go/lib/types"
  19. )
  20. type (
  21. Date = types.Date
  22. DateTime = types.DateTime
  23. UUID = types.UUID
  24. )
  25. var (
  26. ErrInsertInNotBatchMode = errors.New("insert statement supported only in the batch mode (use begin/commit)")
  27. ErrLimitDataRequestInTx = errors.New("data request has already been prepared in transaction")
  28. )
  29. var (
  30. splitInsertRe = regexp.MustCompile(`(?i)\sVALUES\s*\(`)
  31. )
  32. type logger func(format string, v ...interface{})
  33. type clickhouse struct {
  34. sync.Mutex
  35. data.ServerInfo
  36. data.ClientInfo
  37. logf logger
  38. conn *connect
  39. block *data.Block
  40. buffer *bufio.Writer
  41. decoder *binary.Decoder
  42. encoder *binary.Encoder
  43. settings *querySettings
  44. compress bool
  45. blockSize int
  46. inTransaction bool
  47. }
  48. func (ch *clickhouse) Prepare(query string) (driver.Stmt, error) {
  49. return ch.prepareContext(context.Background(), query)
  50. }
  51. func (ch *clickhouse) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
  52. return ch.prepareContext(ctx, query)
  53. }
  54. func (ch *clickhouse) prepareContext(ctx context.Context, query string) (driver.Stmt, error) {
  55. ch.logf("[prepare] %s", query)
  56. switch {
  57. case ch.conn.closed:
  58. return nil, driver.ErrBadConn
  59. case ch.block != nil:
  60. return nil, ErrLimitDataRequestInTx
  61. case isInsert(query):
  62. if !ch.inTransaction {
  63. return nil, ErrInsertInNotBatchMode
  64. }
  65. return ch.insert(query)
  66. }
  67. return &stmt{
  68. ch: ch,
  69. query: query,
  70. numInput: numInput(query),
  71. }, nil
  72. }
  73. func (ch *clickhouse) insert(query string) (_ driver.Stmt, err error) {
  74. if err := ch.sendQuery(splitInsertRe.Split(query, -1)[0] + " VALUES "); err != nil {
  75. return nil, err
  76. }
  77. if ch.block, err = ch.readMeta(); err != nil {
  78. return nil, err
  79. }
  80. return &stmt{
  81. ch: ch,
  82. isInsert: true,
  83. }, nil
  84. }
  85. func (ch *clickhouse) Begin() (driver.Tx, error) {
  86. return ch.beginTx(context.Background(), txOptions{})
  87. }
  88. func (ch *clickhouse) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
  89. return ch.beginTx(ctx, txOptions{
  90. Isolation: int(opts.Isolation),
  91. ReadOnly: opts.ReadOnly,
  92. })
  93. }
  94. type txOptions struct {
  95. Isolation int
  96. ReadOnly bool
  97. }
  98. func (ch *clickhouse) beginTx(ctx context.Context, opts txOptions) (*clickhouse, error) {
  99. ch.logf("[begin] tx=%t, data=%t", ch.inTransaction, ch.block != nil)
  100. switch {
  101. case ch.inTransaction:
  102. return nil, sql.ErrTxDone
  103. case ch.conn.closed:
  104. return nil, driver.ErrBadConn
  105. }
  106. if finish := ch.watchCancel(ctx); finish != nil {
  107. defer finish()
  108. }
  109. ch.block = nil
  110. ch.inTransaction = true
  111. return ch, nil
  112. }
  113. func (ch *clickhouse) Commit() error {
  114. ch.logf("[commit] tx=%t, data=%t", ch.inTransaction, ch.block != nil)
  115. defer func() {
  116. if ch.block != nil {
  117. ch.block.Reset()
  118. ch.block = nil
  119. }
  120. ch.inTransaction = false
  121. }()
  122. switch {
  123. case !ch.inTransaction:
  124. return sql.ErrTxDone
  125. case ch.conn.closed:
  126. return driver.ErrBadConn
  127. }
  128. if ch.block != nil {
  129. if err := ch.writeBlock(ch.block); err != nil {
  130. return err
  131. }
  132. // Send empty block as marker of end of data.
  133. if err := ch.writeBlock(&data.Block{}); err != nil {
  134. return err
  135. }
  136. if err := ch.encoder.Flush(); err != nil {
  137. return err
  138. }
  139. return ch.process()
  140. }
  141. return nil
  142. }
  143. func (ch *clickhouse) Rollback() error {
  144. ch.logf("[rollback] tx=%t, data=%t", ch.inTransaction, ch.block != nil)
  145. if !ch.inTransaction {
  146. return sql.ErrTxDone
  147. }
  148. if ch.block != nil {
  149. ch.block.Reset()
  150. }
  151. ch.block = nil
  152. ch.buffer = nil
  153. ch.inTransaction = false
  154. return ch.conn.Close()
  155. }
  156. func (ch *clickhouse) CheckNamedValue(nv *driver.NamedValue) error {
  157. switch nv.Value.(type) {
  158. case column.IP, column.UUID:
  159. return nil
  160. case nil, []byte, int8, int16, int32, int64, uint8, uint16, uint32, uint64, float32, float64, string, time.Time:
  161. return nil
  162. }
  163. switch v := nv.Value.(type) {
  164. case
  165. []int, []int8, []int16, []int32, []int64,
  166. []uint, []uint8, []uint16, []uint32, []uint64,
  167. []float32, []float64,
  168. []string:
  169. return nil
  170. case net.IP, *net.IP:
  171. return nil
  172. case driver.Valuer:
  173. value, err := v.Value()
  174. if err != nil {
  175. return err
  176. }
  177. nv.Value = value
  178. default:
  179. switch value := reflect.ValueOf(nv.Value); value.Kind() {
  180. case reflect.Slice:
  181. return nil
  182. case reflect.Bool:
  183. nv.Value = uint8(0)
  184. if value.Bool() {
  185. nv.Value = uint8(1)
  186. }
  187. case reflect.Int8:
  188. nv.Value = int8(value.Int())
  189. case reflect.Int16:
  190. nv.Value = int16(value.Int())
  191. case reflect.Int32:
  192. nv.Value = int32(value.Int())
  193. case reflect.Int64:
  194. nv.Value = value.Int()
  195. case reflect.Uint8:
  196. nv.Value = uint8(value.Uint())
  197. case reflect.Uint16:
  198. nv.Value = uint16(value.Uint())
  199. case reflect.Uint32:
  200. nv.Value = uint32(value.Uint())
  201. case reflect.Uint64:
  202. nv.Value = uint64(value.Uint())
  203. case reflect.Float32:
  204. nv.Value = float32(value.Float())
  205. case reflect.Float64:
  206. nv.Value = float64(value.Float())
  207. case reflect.String:
  208. nv.Value = value.String()
  209. }
  210. }
  211. return nil
  212. }
  213. func (ch *clickhouse) Close() error {
  214. ch.block = nil
  215. return ch.conn.Close()
  216. }
  217. func (ch *clickhouse) process() error {
  218. packet, err := ch.decoder.Uvarint()
  219. if err != nil {
  220. return err
  221. }
  222. for {
  223. switch packet {
  224. case protocol.ServerPong:
  225. ch.logf("[process] <- pong")
  226. return nil
  227. case protocol.ServerException:
  228. ch.logf("[process] <- exception")
  229. return ch.exception()
  230. case protocol.ServerProgress:
  231. progress, err := ch.progress()
  232. if err != nil {
  233. return err
  234. }
  235. ch.logf("[process] <- progress: rows=%d, bytes=%d, total rows=%d",
  236. progress.rows,
  237. progress.bytes,
  238. progress.totalRows,
  239. )
  240. case protocol.ServerProfileInfo:
  241. profileInfo, err := ch.profileInfo()
  242. if err != nil {
  243. return err
  244. }
  245. ch.logf("[process] <- profiling: rows=%d, bytes=%d, blocks=%d", profileInfo.rows, profileInfo.bytes, profileInfo.blocks)
  246. case protocol.ServerData:
  247. block, err := ch.readBlock()
  248. if err != nil {
  249. return err
  250. }
  251. ch.logf("[process] <- data: packet=%d, columns=%d, rows=%d", packet, block.NumColumns, block.NumRows)
  252. case protocol.ServerEndOfStream:
  253. ch.logf("[process] <- end of stream")
  254. return nil
  255. default:
  256. ch.conn.Close()
  257. return fmt.Errorf("[process] unexpected packet [%d] from server", packet)
  258. }
  259. if packet, err = ch.decoder.Uvarint(); err != nil {
  260. return err
  261. }
  262. }
  263. }
  264. func (ch *clickhouse) cancel() error {
  265. ch.logf("[cancel request]")
  266. // even if we fail to write the cancel, we still need to close
  267. err := ch.encoder.Uvarint(protocol.ClientCancel)
  268. if err == nil {
  269. err = ch.encoder.Flush()
  270. }
  271. // return the close error if there was one, otherwise return the write error
  272. if cerr := ch.conn.Close(); cerr != nil {
  273. return cerr
  274. }
  275. return err
  276. }
  277. func (ch *clickhouse) watchCancel(ctx context.Context) func() {
  278. if done := ctx.Done(); done != nil {
  279. finished := make(chan struct{})
  280. go func() {
  281. select {
  282. case <-done:
  283. ch.cancel()
  284. finished <- struct{}{}
  285. ch.logf("[cancel] <- done")
  286. case <-finished:
  287. ch.logf("[cancel] <- finished")
  288. }
  289. }()
  290. return func() {
  291. select {
  292. case <-finished:
  293. case finished <- struct{}{}:
  294. }
  295. }
  296. }
  297. return func() {}
  298. }
  299. func (ch *clickhouse) ExecContext(ctx context.Context, query string,
  300. args []driver.NamedValue) (driver.Result, error) {
  301. finish := ch.watchCancel(ctx)
  302. defer finish()
  303. stmt, err := ch.PrepareContext(ctx, query)
  304. if err != nil {
  305. return nil, err
  306. }
  307. dargs := make([]driver.Value, len(args))
  308. for i, nv := range args {
  309. dargs[i] = nv.Value
  310. }
  311. return stmt.Exec(dargs)
  312. }