Make purpose part of the pool key #4986

ref DEV-2444
authgear · Jan 20, 2025 · 55361dd · 55361dd
2 parents 70dcb0f + 93ed942
commit 55361dd
Show file tree

Hide file tree

Showing 13 changed files with 200 additions and 38 deletions.
diff --git a/cmd/authgear/cmd/cmdaudit/audit.go b/cmd/authgear/cmd/cmdaudit/audit.go
@@ -9,6 +9,7 @@ import (
 
 	authgearcmd "github.com/authgear/authgear-server/cmd/authgear/cmd"
 	dbutil "github.com/authgear/authgear-server/pkg/lib/db/util"
+	"github.com/authgear/authgear-server/pkg/lib/infra/db"
 	"github.com/authgear/authgear-server/pkg/util/sqlmigrate"
 )
 
@@ -227,7 +228,10 @@ var cmdAuditDatabaseDump = &cobra.Command{
 		}
 
 		dumper := dbutil.NewDumper(
-			dbURL,
+			db.ConnectionInfo{
+				Purpose:     db.ConnectionPurposeAuditReadOnly,
+				DatabaseURL: dbURL,
+			},
 			dbSchema,
 			outputDir,
 			args,
@@ -257,7 +261,10 @@ var cmdAuditDatabaseRestore = &cobra.Command{
 		}
 
 		restorer := dbutil.NewRestorer(
-			dbURL,
+			db.ConnectionInfo{
+				Purpose:     db.ConnectionPurposeAuditReadWrite,
+				DatabaseURL: dbURL,
+			},
 			dbSchema,
 			inputDir,
 			args,

diff --git a/cmd/authgear/cmd/cmddatabase/database.go b/cmd/authgear/cmd/cmddatabase/database.go
@@ -10,6 +10,7 @@ import (
 
 	authgearcmd "github.com/authgear/authgear-server/cmd/authgear/cmd"
 	dbutil "github.com/authgear/authgear-server/pkg/lib/db/util"
+	"github.com/authgear/authgear-server/pkg/lib/infra/db"
 	"github.com/authgear/authgear-server/pkg/util/sqlmigrate"
 )
 
@@ -193,7 +194,10 @@ var cmdDump = &cobra.Command{
 		}
 
 		dumper := dbutil.NewDumper(
-			dbURL,
+			db.ConnectionInfo{
+				Purpose:     db.ConnectionPurposeApp,
+				DatabaseURL: dbURL,
+			},
 			dbSchema,
 			outputDir,
 			args,
@@ -223,7 +227,10 @@ var cmdRestore = &cobra.Command{
 		}
 
 		restorer := dbutil.NewRestorer(
-			dbURL,
+			db.ConnectionInfo{
+				Purpose:     db.ConnectionPurposeApp,
+				DatabaseURL: dbURL,
+			},
 			dbSchema,
 			inputDir,
 			args,

diff --git a/cmd/portal/cmd/cmddatabase/database.go b/cmd/portal/cmd/cmddatabase/database.go
@@ -9,6 +9,7 @@ import (
 
 	portalcmd "github.com/authgear/authgear-server/cmd/portal/cmd"
 	dbutil "github.com/authgear/authgear-server/pkg/lib/db/util"
+	"github.com/authgear/authgear-server/pkg/lib/infra/db"
 	"github.com/authgear/authgear-server/pkg/util/sqlmigrate"
 )
 
@@ -192,7 +193,10 @@ var cmdDump = &cobra.Command{
 		}
 
 		dumper := dbutil.NewDumper(
-			dbURL,
+			db.ConnectionInfo{
+				Purpose:     db.ConnectionPurposeGlobal,
+				DatabaseURL: dbURL,
+			},
 			dbSchema,
 			outputDir,
 			args,
@@ -222,7 +226,10 @@ var cmdRestore = &cobra.Command{
 		}
 
 		restorer := dbutil.NewRestorer(
-			dbURL,
+			db.ConnectionInfo{
+				Purpose:     db.ConnectionPurposeGlobal,
+				DatabaseURL: dbURL,
+			},
 			dbSchema,
 			inputDir,
 			args,

diff --git a/pkg/lib/db/util/dump.go b/pkg/lib/db/util/dump.go
@@ -16,7 +16,7 @@ import (
 )
 
 type Dumper struct {
-	DatabaseURL    string
+	ConnectionInfo db.ConnectionInfo
 	DatabaseSchema string
 	OutputDir      string
 	AppIDs         []string
@@ -29,7 +29,7 @@ type Dumper struct {
 }
 
 func NewDumper(
-	databaseURL string,
+	connectionInfo db.ConnectionInfo,
 	databaseSchema string,
 	outputDir string,
 	appIDs []string,
@@ -42,8 +42,8 @@ func NewDumper(
 	pool := db.NewPool()
 	handle := db.NewHookHandle(
 		pool,
+		connectionInfo,
 		db.ConnectionOptions{
-			DatabaseURL:           databaseURL,
 			MaxOpenConnection:     1,
 			MaxIdleConnection:     1,
 			MaxConnectionLifetime: 1800 * time.Second,
@@ -54,7 +54,7 @@ func NewDumper(
 	sqlExecutor := &db.SQLExecutor{}
 	sqlBuilder := db.NewSQLBuilder(databaseSchema)
 	return &Dumper{
-		DatabaseURL:    databaseURL,
+		ConnectionInfo: connectionInfo,
 		DatabaseSchema: databaseSchema,
 		OutputDir:      outputDir,
 		AppIDs:         appIDs,

diff --git a/pkg/lib/db/util/restore.go b/pkg/lib/db/util/restore.go
@@ -13,7 +13,7 @@ import (
 )
 
 type Restorer struct {
-	DatabaseURL    string
+	ConnectionInfo db.ConnectionInfo
 	DatabaseSchema string
 	InputDir       string
 	AppIDs         []string
@@ -26,7 +26,7 @@ type Restorer struct {
 }
 
 func NewRestorer(
-	databaseURL string,
+	connectionInfo db.ConnectionInfo,
 	databaseSchema string,
 	inputDir string,
 	appIDs []string,
@@ -39,8 +39,8 @@ func NewRestorer(
 	pool := db.NewPool()
 	handle := db.NewHookHandle(
 		pool,
+		connectionInfo,
 		db.ConnectionOptions{
-			DatabaseURL:           databaseURL,
 			MaxOpenConnection:     1,
 			MaxIdleConnection:     1,
 			MaxConnectionLifetime: 1800 * time.Second,
@@ -51,7 +51,7 @@ func NewRestorer(
 	sqlExecutor := &db.SQLExecutor{}
 	sqlBuilder := db.NewSQLBuilder(databaseSchema)
 	return &Restorer{
-		DatabaseURL:    databaseURL,
+		ConnectionInfo: connectionInfo,
 		DatabaseSchema: databaseSchema,
 		InputDir:       inputDir,
 		AppIDs:         appIDs,

diff --git a/pkg/lib/infra/db/appdb/deps.go b/pkg/lib/infra/db/appdb/deps.go
@@ -73,14 +73,17 @@ func NewHandle(
 	credentials *config.DatabaseCredentials,
 	lf *log.Factory,
 ) *Handle {
+	info := db.ConnectionInfo{
+		Purpose:     db.ConnectionPurposeApp,
+		DatabaseURL: credentials.DatabaseURL,
+	}
 	opts := db.ConnectionOptions{
-		DatabaseURL:           credentials.DatabaseURL,
 		MaxOpenConnection:     cfg.MaxOpenConn,
 		MaxIdleConnection:     cfg.MaxIdleConn,
 		MaxConnectionLifetime: cfg.ConnMaxLifetimeSeconds.Duration(),
 		IdleConnectionTimeout: cfg.ConnMaxIdleTimeSeconds.Duration(),
 	}
 	return &Handle{
-		db.NewHookHandle(pool, opts, lf),
+		db.NewHookHandle(pool, info, opts, lf),
 	}
 }
diff --git a/pkg/lib/infra/db/auditdb/deps.go b/pkg/lib/infra/db/auditdb/deps.go
@@ -102,15 +102,19 @@ func NewReadHandle(
 		return nil
 	}
 
+	info := db.ConnectionInfo{
+		Purpose:     db.ConnectionPurposeAuditReadOnly,
+		DatabaseURL: credentials.DatabaseURL,
+	}
+
 	opts := db.ConnectionOptions{
-		DatabaseURL:           credentials.DatabaseURL,
 		MaxOpenConnection:     cfg.MaxOpenConn,
 		MaxIdleConnection:     cfg.MaxIdleConn,
 		MaxConnectionLifetime: cfg.ConnMaxLifetimeSeconds.Duration(),
 		IdleConnectionTimeout: cfg.ConnMaxIdleTimeSeconds.Duration(),
 	}
 	return &ReadHandle{
-		db.NewHookHandle(pool, opts, lf),
+		db.NewHookHandle(pool, info, opts, lf),
 	}
 }
 
@@ -132,15 +136,19 @@ func NewWriteHandle(
 		return nil
 	}
 
+	info := db.ConnectionInfo{
+		Purpose:     db.ConnectionPurposeAuditReadWrite,
+		DatabaseURL: credentials.DatabaseURL,
+	}
+
 	opts := db.ConnectionOptions{
-		DatabaseURL:           credentials.DatabaseURL,
 		MaxOpenConnection:     cfg.MaxOpenConn,
 		MaxIdleConnection:     cfg.MaxIdleConn,
 		MaxConnectionLifetime: cfg.ConnMaxLifetimeSeconds.Duration(),
 		IdleConnectionTimeout: cfg.ConnMaxIdleTimeSeconds.Duration(),
 	}
 	return &WriteHandle{
-		db.NewHookHandle(pool, opts, lf),
+		db.NewHookHandle(pool, info, opts, lf),
 	}
 }
 

diff --git a/pkg/lib/infra/db/connection_options.go b/pkg/lib/infra/db/connection_options.go
@@ -4,8 +4,22 @@ import (
 	"time"
 )
 
+type ConnectionPurpose string
+
+const (
+	ConnectionPurposeGlobal         = "global"
+	ConnectionPurposeApp            = "app"
+	ConnectionPurposeAuditReadOnly  = "audit_read_only"
+	ConnectionPurposeAuditReadWrite = "audit_read_write"
+	ConnectionPurposeSearch         = "search"
+)
+
+type ConnectionInfo struct {
+	Purpose     ConnectionPurpose
+	DatabaseURL string
+}
+
 type ConnectionOptions struct {
-	DatabaseURL           string
 	MaxOpenConnection     int
 	MaxIdleConnection     int
 	MaxConnectionLifetime time.Duration

diff --git a/pkg/lib/infra/db/globaldb/deps.go b/pkg/lib/infra/db/globaldb/deps.go
@@ -48,14 +48,18 @@ func NewHandle(
 	cfg *config.DatabaseEnvironmentConfig,
 	lf *log.Factory,
 ) *Handle {
+	info := db.ConnectionInfo{
+		Purpose:     db.ConnectionPurposeGlobal,
+		DatabaseURL: credentials.DatabaseURL,
+	}
+
 	opts := db.ConnectionOptions{
-		DatabaseURL:           credentials.DatabaseURL,
 		MaxOpenConnection:     cfg.MaxOpenConn,
 		MaxIdleConnection:     cfg.MaxIdleConn,
 		MaxConnectionLifetime: time.Second * time.Duration(cfg.ConnMaxLifetimeSeconds),
 		IdleConnectionTimeout: time.Second * time.Duration(cfg.ConnMaxIdleTimeSeconds),
 	}
 	return &Handle{
-		db.NewHookHandle(pool, opts, lf),
+		db.NewHookHandle(pool, info, opts, lf),
 	}
 }
diff --git a/pkg/lib/infra/db/hook_handle.go b/pkg/lib/infra/db/hook_handle.go
@@ -21,6 +21,7 @@ type hookHandleContextValue struct {
 
 type HookHandle struct {
 	Pool              *Pool
+	ConnectionInfo    ConnectionInfo
 	ConnectionOptions ConnectionOptions
 	Logger            *log.Logger
 }
@@ -51,9 +52,10 @@ func mustGetTxLike(ctx context.Context) txLike {
 
 var _ Handle = (*HookHandle)(nil)
 
-func NewHookHandle(pool *Pool, opts ConnectionOptions, lf *log.Factory) *HookHandle {
+func NewHookHandle(pool *Pool, info ConnectionInfo, opts ConnectionOptions, lf *log.Factory) *HookHandle {
 	return &HookHandle{
 		Pool:              pool,
+		ConnectionInfo:    info,
 		ConnectionOptions: opts,
 		Logger:            lf.New("db-handle"),
 	}
@@ -245,13 +247,14 @@ func rollbackTx(logger *log.Logger, tx *sql.Tx) error {
 
 func (h *HookHandle) openDB() (*sql.DB, error) {
 	h.Logger.WithFields(map[string]interface{}{
+		"purpose":                    h.ConnectionInfo.Purpose,
 		"max_open_conns":             h.ConnectionOptions.MaxOpenConnection,
 		"max_idle_conns":             h.ConnectionOptions.MaxIdleConnection,
 		"conn_max_lifetime_seconds":  h.ConnectionOptions.MaxConnectionLifetime.Seconds(),
 		"conn_max_idle_time_seconds": h.ConnectionOptions.IdleConnectionTimeout.Seconds(),
 	}).Debug("open database")
 
-	db, err := h.Pool.Open(h.ConnectionOptions)
+	db, err := h.Pool.Open(h.ConnectionInfo, h.ConnectionOptions)
 	if err != nil {
 		return nil, fmt.Errorf("hook-handle: failed to connect to database: %w", err)
 	}

diff --git a/pkg/lib/infra/db/pool.go b/pkg/lib/infra/db/pool.go
@@ -8,38 +8,38 @@ import (
 	"github.com/authgear/authgear-server/pkg/util/otelutil"
 )
 
+var actualPoolOpener = openPostgresDB
+
 type Pool struct {
 	closed     bool
 	closeMutex sync.RWMutex
 
-	cache      map[string]*sql.DB
+	cache      map[ConnectionInfo]*sql.DB
 	cacheMutex sync.RWMutex
 }
 
 func NewPool() *Pool {
-	return &Pool{cache: map[string]*sql.DB{}}
+	return &Pool{cache: map[ConnectionInfo]*sql.DB{}}
 }
 
-func (p *Pool) Open(opts ConnectionOptions) (db *sql.DB, err error) {
-	source := opts.DatabaseURL
-
+func (p *Pool) Open(info ConnectionInfo, opts ConnectionOptions) (db *sql.DB, err error) {
 	p.closeMutex.RLock()
 	defer func() { p.closeMutex.RUnlock() }()
 	if p.closed {
 		return nil, errors.New("db: pool is closed")
 	}
 
 	p.cacheMutex.RLock()
-	db, exists := p.cache[source]
+	db, exists := p.cache[info]
 	p.cacheMutex.RUnlock()
 
 	if !exists {
 		p.cacheMutex.Lock()
-		db, exists = p.cache[source]
+		db, exists = p.cache[info]
 		if !exists {
-			db, err = p.openPostgresDB(opts)
+			db, err = actualPoolOpener(info, opts)
 			if err == nil {
-				p.cache[source] = db
+				p.cache[info] = db
 			}
 		}
 		p.cacheMutex.Unlock()
@@ -65,8 +65,8 @@ func (p *Pool) Close() (err error) {
 	return
 }
 
-func (p *Pool) openPostgresDB(opts ConnectionOptions) (*sql.DB, error) {
-	pgdb, err := otelutil.OTelSQLOpenPostgres(opts.DatabaseURL)
+func openPostgresDB(info ConnectionInfo, opts ConnectionOptions) (*sql.DB, error) {
+	pgdb, err := otelutil.OTelSQLOpenPostgres(info.DatabaseURL)
 	if err != nil {
 		return nil, err
 	}