feat: different db context, partition for cache (#13)

cache/cache.go

@@ -39,19 +39,21 @@ var DefaultConfig = Config{
 	CleanupInterval: 15 * time.Second,
 }
 
-func WithTTL(ttl time.Duration) func(*Config) {
+type Option func(*Config)
+
+func WithTTL(ttl time.Duration) Option {
 	return func(config *Config) {
 		config.TTL = ttl
 	}
 }
 
-func WithCleanupInterval(val time.Duration) func(*Config) {
+func WithCleanupInterval(val time.Duration) Option {
 	return func(config *Config) {
 		config.CleanupInterval = val
 	}
 }
 
-func WithCapacity(c int) func(*Config) {
+func WithCapacity(c int) Option {
 	return func(config *Config) {
 		config.Capacity = c
 	}
@@ -84,7 +86,7 @@ type DeleteHandler func(*valuesItem) *keyValue
 type EvictionHandler func(string, interface{})
 
 // Init initializes cache struct fields and starts janitor process.
-func (c *Cache) Init(deleteHandler DeleteHandler, opts ...func(*Config)) {
+func (c *Cache) Init(deleteHandler DeleteHandler, opts ...Option) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 

cache/lru.go

@@ -12,7 +12,7 @@ type LRUCache struct {
 
 // NewLRUCache creates and initializes a new cache object.
 // This one is based on LRU KV with TTL
-func NewLRUCache(autoRefresh bool, opts ...func(*Config)) *LRUCache {
+func NewLRUCache(autoRefresh bool, opts ...Option) *LRUCache {
 	cache := LRUCache{
 		autoRefresh: autoRefresh,
 	}

cache/lru_set.go

@@ -11,7 +11,7 @@ type LRUSetCache struct {
 
 // NewLRUSetCache creates and initializes a new cache object.
 // This one is based on LRU K->List with TTL
-func NewLRUSetCache(opts ...func(*Config)) *LRUSetCache {
+func NewLRUSetCache(opts ...Option) *LRUSetCache {
 	cache := LRUSetCache{}
 	cache.Cache.Init(cache.deleteHandler, opts...)
 

database/interface.go

@@ -1,11 +1,15 @@
 package database
 
 import (
+	"context"
+
 	"github.com/go-pg/pg/v9"
 )
 
 //go:generate go run github.com/vektra/mockery/cmd/mockery -name DBContext
 type DBContext interface {
+	Context() context.Context
+	Unwrap() interface{}
 	Schema() string
 }
 

database/manager/live_manager.go

@@ -1,8 +1,6 @@
 package manager
 
 import (
-	"context"
-
 	"github.com/go-pg/pg/v9/orm"
 
 	"github.com/Syncano/pkg-go/v2/database"
@@ -25,39 +23,23 @@ func NewLiveTenantManager(c database.DBContext, db *database.DB) *LiveManager {
 
 // Query returns only alive objects.
 func (m *LiveManager) Query(o interface{}) *orm.Query {
-	return m.QueryContext(context.Background(), o)
-}
-
-func (m *LiveManager) QueryContext(ctx context.Context, o interface{}) *orm.Query {
-	return m.DB().ModelContext(ctx, o).Where("?TableAlias._is_live IS TRUE")
+	return m.DB().ModelContext(m.dbCtx.Context(), o).Where("?TableAlias._is_live IS TRUE")
 }
 
 // All returns all objects, irrelevant if they are alive or not.
 func (m *LiveManager) All(o interface{}) *orm.Query {
-	return m.AllContext(context.Background(), o)
-}
-
-func (m *LiveManager) AllContext(ctx context.Context, o interface{}) *orm.Query {
-	return m.DB().ModelContext(ctx, o)
+	return m.DB().ModelContext(m.dbCtx.Context(), o)
 }
 
 // Dead returns dead objects.
 func (m *LiveManager) Dead(o interface{}) *orm.Query {
-	return m.DeadContext(context.Background(), o)
-}
-
-func (m *LiveManager) DeadContext(ctx context.Context, o interface{}) *orm.Query {
-	return m.DB().ModelContext(ctx, o).Where("?TableAlias._is_live IS NULL")
+	return m.DB().ModelContext(m.dbCtx.Context(), o).Where("?TableAlias._is_live IS NULL")
 }
 
 // Delete is a soft delete for live objects.
 func (m *LiveManager) Delete(model interface{}) error {
-	return m.DeleteContext(context.Background(), model)
-}
-
-func (m *LiveManager) DeleteContext(ctx context.Context, model interface{}) error {
 	db := m.DB()
-	if _, err := db.ModelContext(ctx, model).WherePK().Set("_is_live = ?", false).Update(); err != nil {
+	if _, err := db.ModelContext(m.dbCtx.Context(), model).WherePK().Set("_is_live = ?", false).Update(); err != nil {
 		return err
 	}
 

database/manager/manager.go

@@ -59,21 +59,13 @@ func (m *Manager) SetDB(db orm.DB) {
 
 // Query returns all objects.
 func (m *Manager) Query(o interface{}) *orm.Query {
-	return m.QueryContext(context.Background(), o)
-}
-
-func (m *Manager) QueryContext(ctx context.Context, o interface{}) *orm.Query {
-	return m.DB().ModelContext(ctx, o)
+	return m.DB().ModelContext(m.dbCtx.Context(), o)
 }
 
 // Insert creates object.
 func (m *Manager) Insert(model interface{}) error {
-	return m.InsertContext(context.Background(), model)
-}
-
-func (m *Manager) InsertContext(ctx context.Context, model interface{}) error {
 	db := m.DB()
-	if _, err := db.ModelContext(ctx, model).Insert(model); err != nil {
+	if _, err := db.ModelContext(m.dbCtx.Context(), model).Insert(model); err != nil {
 		return err
 	}
 
@@ -82,12 +74,8 @@ func (m *Manager) InsertContext(ctx context.Context, model interface{}) error {
 
 // Update updates object.
 func (m *Manager) Update(model interface{}, fields ...string) error {
-	return m.UpdateContext(context.Background(), model, fields...)
-}
-
-func (m *Manager) UpdateContext(ctx context.Context, model interface{}, fields ...string) error {
 	db := m.DB()
-	if _, err := db.ModelContext(ctx, model).Column(fields...).WherePK().Update(); err != nil {
+	if _, err := db.ModelContext(m.dbCtx.Context(), model).Column(fields...).WherePK().Update(); err != nil {
 		return err
 	}
 

database/wrapper.go

@@ -1,27 +1,41 @@
 package database
 
-import "context"
+import (
+	"context"
+)
 
 type wrappedContext struct {
-	context.Context
+	wrapped      interface{}
+	ctx          context.Context
 	schemaGetter func() string
 }
 
+var _ DBContext = (*wrappedContext)(nil)
+
 func (w *wrappedContext) Schema() string {
 	return w.schemaGetter()
 }
 
-func WrapContext(ctx context.Context) DBContext {
-	return WrapContextWithSchema(ctx, "public")
+func (w *wrappedContext) Context() context.Context {
+	return w.ctx
+}
+
+func (w *wrappedContext) Unwrap() interface{} {
+	return w.wrapped
+}
+
+func WrapContext(ctx context.Context, data interface{}) DBContext {
+	return WrapContextWithSchema(ctx, "public", data)
 }
 
-func WrapContextWithSchema(ctx context.Context, schema string) DBContext {
-	return WrapContextWithSchemaGetter(ctx, func() string { return schema })
+func WrapContextWithSchema(ctx context.Context, schema string, data interface{}) DBContext {
+	return WrapContextWithSchemaGetter(ctx, func() string { return schema }, data)
 }
 
-func WrapContextWithSchemaGetter(ctx context.Context, schemaGetter func() string) DBContext {
+func WrapContextWithSchemaGetter(ctx context.Context, schemaGetter func() string, data interface{}) DBContext {
 	return &wrappedContext{
-		Context:      ctx,
+		ctx:          ctx,
 		schemaGetter: schemaGetter,
+		wrapped:      data,
 	}
 }

go.mod

@@ -4,12 +4,12 @@ go 1.14
 
 require (
 	cloud.google.com/go/storage v1.10.0
-	github.com/TheZeroSlave/zapsentry v1.4.1
+	github.com/TheZeroSlave/zapsentry v1.5.0
 	github.com/alexandrevicenzi/unchained v1.3.0
-	github.com/aws/aws-sdk-go v1.33.18
+	github.com/aws/aws-sdk-go v1.34.7
 	github.com/cloudfoundry/gosigar v1.1.0
 	github.com/getsentry/sentry-go v0.7.0
-	github.com/go-pg/pg/v9 v9.1.6
+	github.com/go-pg/pg/v9 v9.1.7
 	github.com/go-redis/cache/v7 v7.0.2
 	github.com/go-redis/redis/v7 v7.4.0
 	github.com/google/uuid v1.1.1 // indirect
@@ -34,7 +34,7 @@ require (
 	go.opencensus.io v0.22.4
 	go.uber.org/zap v1.15.0
 	golang.org/x/crypto v0.0.0-20200728195943-123391ffb6de
-	golang.org/x/tools v0.0.0-20200804011535-6c149bb5ef0d
-	google.golang.org/api v0.29.0
+	golang.org/x/tools v0.0.0-20200819140908-cf83efe03cf8
+	google.golang.org/api v0.30.0
 	google.golang.org/grpc v1.31.0
 )