improve clickhouse driver

This commit is contained in:
John Guo
2022-04-12 21:31:51 +08:00
parent 7785082f19
commit ee29b28575
8 changed files with 59 additions and 43 deletions

View File

@ -53,7 +53,7 @@ func (c *Core) Ctx(ctx context.Context) DB {
panic(err)
}
newCore.ctx = WithDB(ctx, newCore.db)
newCore.ctx = c.injectInternalCtxData(newCore.ctx)
newCore.ctx = c.InjectInternalCtxData(newCore.ctx)
return newCore.db
}
@ -64,7 +64,7 @@ func (c *Core) GetCtx() context.Context {
if ctx == nil {
ctx = context.TODO()
}
return c.injectInternalCtxData(ctx)
return c.InjectInternalCtxData(ctx)
}
// GetCtxTimeout returns the context and cancel function for specified timeout type.

View File

@ -31,10 +31,10 @@ const (
// for example: `clickhouse`. The `clickhouse` does not support fetching insert/update results,
// but returns errors when execute `RowsAffected`. It here ignores the calling of `RowsAffected`
// to avoid triggering errors, rather than ignoring errors after they are triggered.
ignoreResultInCtx gctx.StrKey = "IgnoreResult"
ignoreResultKeyInCtx gctx.StrKey = "IgnoreResult"
)
func (c *Core) injectInternalCtxData(ctx context.Context) context.Context {
func (c *Core) InjectInternalCtxData(ctx context.Context) context.Context {
// If the internal data is already injected, it does nothing.
if ctx.Value(internalCtxDataKeyInCtx) != nil {
return ctx
@ -44,25 +44,23 @@ func (c *Core) injectInternalCtxData(ctx context.Context) context.Context {
})
}
func (c *Core) InjectIgnoreResult(ctx context.Context) context.Context {
if ctx.Value(ignoreResultInCtx) != nil {
return ctx
}
return context.WithValue(ctx, ignoreResultInCtx, &internalCtxData{
DB: c.db,
})
}
func (c *Core) getIgnoreResultFromCtx(ctx context.Context) *internalCtxData {
if v := ctx.Value(ignoreResultInCtx); v != nil {
return v.(*internalCtxData)
}
return nil
}
func (c *Core) getInternalCtxDataFromCtx(ctx context.Context) *internalCtxData {
func (c *Core) GetInternalCtxDataFromCtx(ctx context.Context) *internalCtxData {
if v := ctx.Value(internalCtxDataKeyInCtx); v != nil {
return v.(*internalCtxData)
}
return nil
}
func (c *Core) InjectIgnoreResult(ctx context.Context) context.Context {
if ctx.Value(ignoreResultKeyInCtx) != nil {
return ctx
}
return context.WithValue(ctx, ignoreResultKeyInCtx, true)
}
func (c *Core) GetIgnoreResultFromCtx(ctx context.Context) bool {
if ctx.Value(ignoreResultKeyInCtx) != nil {
return true
}
return false
}

View File

@ -159,7 +159,7 @@ func (tx *TX) transactionKeyForNestedPoint() string {
func (tx *TX) Ctx(ctx context.Context) *TX {
tx.ctx = ctx
if tx.ctx != nil {
tx.ctx = tx.db.GetCore().injectInternalCtxData(tx.ctx)
tx.ctx = tx.db.GetCore().InjectInternalCtxData(tx.ctx)
}
return tx
}

View File

@ -165,7 +165,7 @@ func (c *Core) sqlParsingHandler(ctx context.Context, in sqlParsingHandlerInput)
// DoCommit commits current sql and arguments to underlying sql driver.
func (c *Core) DoCommit(ctx context.Context, in DoCommitInput) (out DoCommitOutput, err error) {
// Inject internal data into ctx, especially for transaction creating.
ctx = c.injectInternalCtxData(ctx)
ctx = c.InjectInternalCtxData(ctx)
var (
sqlTx *sql.Tx
@ -261,7 +261,7 @@ func (c *Core) DoCommit(ctx context.Context, in DoCommitInput) (out DoCommitOutp
}
// Result handling.
switch {
case sqlResult != nil && c.getIgnoreResultFromCtx(ctx) == nil:
case sqlResult != nil && !c.GetIgnoreResultFromCtx(ctx):
rowsAffected, err = sqlResult.RowsAffected()
out.Result = sqlResult
@ -401,7 +401,7 @@ func (c *Core) RowsToResult(ctx context.Context, rows *sql.Rows) (Result, error)
columnNames[k] = v.Name()
}
if len(columnNames) > 0 {
if internalData := c.getInternalCtxDataFromCtx(ctx); internalData != nil {
if internalData := c.GetInternalCtxDataFromCtx(ctx); internalData != nil {
internalData.FirstResultColumn = columnNames[0]
}
}

View File

@ -75,7 +75,7 @@ func (m *Model) getSelectResultFromCache(ctx context.Context, sql string, args .
)
defer func() {
if cacheItem != nil {
if internalData := m.db.GetCore().getInternalCtxDataFromCtx(ctx); internalData != nil {
if internalData := m.db.GetCore().GetInternalCtxDataFromCtx(ctx); internalData != nil {
if internalData.FirstResultColumn == "" {
internalData.FirstResultColumn = cacheItem.FirstResultColumn
}
@ -114,7 +114,7 @@ func (m *Model) saveSelectResultToCache(ctx context.Context, result Result, sql
var cacheItem = &selectCacheItem{
Result: result,
}
if internalData := m.db.GetCore().getInternalCtxDataFromCtx(ctx); internalData != nil {
if internalData := m.db.GetCore().GetInternalCtxDataFromCtx(ctx); internalData != nil {
cacheItem.FirstResultColumn = internalData.FirstResultColumn
}
if errCache := cacheObj.Set(ctx, cacheKey, cacheItem, m.cacheOption.Duration); errCache != nil {

View File

@ -174,7 +174,7 @@ func (m *Model) Value(fieldsAndWhere ...interface{}) (Value, error) {
if len(all) == 0 {
return gvar.New(nil), nil
}
if internalData := m.db.GetCore().getInternalCtxDataFromCtx(ctx); internalData != nil {
if internalData := m.db.GetCore().GetInternalCtxDataFromCtx(ctx); internalData != nil {
record := all[0]
if v, ok := record[internalData.FirstResultColumn]; ok {
return v, nil
@ -391,7 +391,7 @@ func (m *Model) Count(where ...interface{}) (int, error) {
return 0, err
}
if len(all) > 0 {
if internalData := m.db.GetCore().getInternalCtxDataFromCtx(ctx); internalData != nil {
if internalData := m.db.GetCore().GetInternalCtxDataFromCtx(ctx); internalData != nil {
record := all[0]
if v, ok := record[internalData.FirstResultColumn]; ok {
return v.Int(), nil