From 42897a26848da5b96c19f7ebdebe36eff0409c6a Mon Sep 17 00:00:00 2001 From: Tolga Ozen Date: Sat, 4 May 2024 15:33:19 +0300 Subject: [PATCH] refactor(database): streamline default query execution mode setting --- internal/config/config.go | 8 +++--- internal/factories/database.go | 1 - internal/storage/migration.go | 2 +- pkg/cmd/config.go | 1 - pkg/cmd/flags/serve.go | 7 ------ pkg/cmd/serve.go | 1 - pkg/database/postgres/consts.go | 1 - pkg/database/postgres/options.go | 6 ----- pkg/database/postgres/postgres.go | 42 ++++++++++++++++++++----------- 9 files changed, 32 insertions(+), 37 deletions(-) diff --git a/internal/config/config.go b/internal/config/config.go index 078e9dac1..1474a80f6 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -144,10 +144,9 @@ type ( // Database contains configuration for the database. Database struct { - Engine string `mapstructure:"engine"` // Database engine type (e.g., "postgres" or "memory") - URI string `mapstructure:"uri"` // Database connection URI - AutoMigrate bool `mapstructure:"auto_migrate"` // Whether to enable automatic migration - SimpleMode bool `mapstructure:"simple_mode"` + Engine string `mapstructure:"engine"` // Database engine type (e.g., "postgres" or "memory") + URI string `mapstructure:"uri"` // Database connection URI + AutoMigrate bool `mapstructure:"auto_migrate"` // Whether to enable automatic migration MaxOpenConnections int `mapstructure:"max_open_connections"` // Maximum number of open connections to the database MaxIdleConnections int `mapstructure:"max_idle_connections"` // Maximum number of idle connections to the database MaxConnectionLifetime time.Duration `mapstructure:"max_connection_lifetime"` // Maximum duration a connection can be reused @@ -318,7 +317,6 @@ func DefaultConfig() *Config { Database: Database{ Engine: "memory", AutoMigrate: true, - SimpleMode: false, MaxDataPerWrite: 1000, MaxRetries: 10, WatchBufferSize: 100, diff --git a/internal/factories/database.go b/internal/factories/database.go index 42ebe896a..f06589636 100644 --- a/internal/factories/database.go +++ b/internal/factories/database.go @@ -40,7 +40,6 @@ func DatabaseFactory(conf config.Database) (db database.Database, err error) { PQDatabase.WatchBufferSize(conf.WatchBufferSize), PQDatabase.MaxDataPerWrite(conf.MaxDataPerWrite), PQDatabase.MaxRetries(conf.MaxRetries), - PQDatabase.SimpleMode(conf.SimpleMode), ) if err != nil { return nil, err diff --git a/internal/storage/migration.go b/internal/storage/migration.go index 9d9330bfc..302db6584 100644 --- a/internal/storage/migration.go +++ b/internal/storage/migration.go @@ -30,7 +30,7 @@ func Migrate(conf config.Database) (err error) { case database.POSTGRES.String(): // Create a new Postgres database connection var db *PQDatabase.Postgres - db, err = PQDatabase.New(conf.URI, PQDatabase.SimpleMode(conf.SimpleMode)) + db, err = PQDatabase.New(conf.URI) if err != nil { return err } diff --git a/pkg/cmd/config.go b/pkg/cmd/config.go index ebd5f475a..b4343e37e 100644 --- a/pkg/cmd/config.go +++ b/pkg/cmd/config.go @@ -72,7 +72,6 @@ func NewConfigCommand() *cobra.Command { f.String("database-engine", conf.Database.Engine, "data source. e.g. postgres, memory") f.String("database-uri", conf.Database.URI, "uri of your data source to store relation tuples and schema") f.Bool("database-auto-migrate", conf.Database.AutoMigrate, "auto migrate database tables") - f.Bool("database-simple-mode", conf.Database.SimpleMode, "database connection to operate using pgx.QueryExecModeSimpleProtocol") f.Int("database-max-open-connections", conf.Database.MaxOpenConnections, "maximum number of parallel connections that can be made to the database at any time") f.Int("database-max-idle-connections", conf.Database.MaxIdleConnections, "maximum number of idle connections that can be made to the database at any time") f.Duration("database-max-connection-lifetime", conf.Database.MaxConnectionLifetime, "maximum amount of time a connection may be reused") diff --git a/pkg/cmd/flags/serve.go b/pkg/cmd/flags/serve.go index 0d61902a8..924ac84d6 100644 --- a/pkg/cmd/flags/serve.go +++ b/pkg/cmd/flags/serve.go @@ -353,13 +353,6 @@ func RegisterServeFlags(flags *pflag.FlagSet) { panic(err) } - if err = viper.BindPFlag("database.simple_mode", flags.Lookup("database-simple-mode")); err != nil { - panic(err) - } - if err = viper.BindEnv("database.simple_mode", "PERMIFY_DATABASE_SIMPLE_MODE"); err != nil { - panic(err) - } - if err = viper.BindPFlag("database.max_open_connections", flags.Lookup("database-max-open-connections")); err != nil { panic(err) } diff --git a/pkg/cmd/serve.go b/pkg/cmd/serve.go index 5246390f6..8de2d423a 100644 --- a/pkg/cmd/serve.go +++ b/pkg/cmd/serve.go @@ -103,7 +103,6 @@ func NewServeCommand() *cobra.Command { f.String("database-engine", conf.Database.Engine, "data source. e.g. postgres, memory") f.String("database-uri", conf.Database.URI, "uri of your data source to store relation tuples and schema") f.Bool("database-auto-migrate", conf.Database.AutoMigrate, "auto migrate database tables") - f.Bool("database-simple-mode", conf.Database.SimpleMode, "database connection to operate using pgx.QueryExecModeSimpleProtocol") f.Int("database-max-open-connections", conf.Database.MaxOpenConnections, "maximum number of parallel connections that can be made to the database at any time") f.Int("database-max-idle-connections", conf.Database.MaxIdleConnections, "maximum number of idle connections that can be made to the database at any time") f.Duration("database-max-connection-lifetime", conf.Database.MaxConnectionLifetime, "maximum amount of time a connection may be reused") diff --git a/pkg/database/postgres/consts.go b/pkg/database/postgres/consts.go index 99366bae3..603430507 100644 --- a/pkg/database/postgres/consts.go +++ b/pkg/database/postgres/consts.go @@ -6,5 +6,4 @@ const ( _defaultMaxDataPerWrite = 1000 _defaultMaxRetries = 10 _defaultWatchBufferSize = 100 - _defaultSimpleMode = false ) diff --git a/pkg/database/postgres/options.go b/pkg/database/postgres/options.go index e5d94a8d8..5e25939fc 100644 --- a/pkg/database/postgres/options.go +++ b/pkg/database/postgres/options.go @@ -52,9 +52,3 @@ func MaxRetries(v int) Option { c.maxRetries = v } } - -func SimpleMode(v bool) Option { - return func(c *Postgres) { - c.simpleMode = v - } -} diff --git a/pkg/database/postgres/postgres.go b/pkg/database/postgres/postgres.go index 3bf9e3461..2ba044648 100644 --- a/pkg/database/postgres/postgres.go +++ b/pkg/database/postgres/postgres.go @@ -2,9 +2,11 @@ package postgres import ( "context" + "strings" "time" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" "github.com/Masterminds/squirrel" @@ -24,7 +26,6 @@ type Postgres struct { maxConnectionIdleTime time.Duration maxOpenConnections int maxIdleConnections int - simpleMode bool } // New - Creates new postgresql db instance @@ -35,7 +36,6 @@ func New(uri string, opts ...Option) (*Postgres, error) { maxDataPerWrite: _defaultMaxDataPerWrite, maxRetries: _defaultMaxRetries, watchBufferSize: _defaultWatchBufferSize, - simpleMode: _defaultSimpleMode, } // Custom options @@ -55,10 +55,8 @@ func New(uri string, opts ...Option) (*Postgres, error) { return nil, err } - if pg.simpleMode { - writeConfig.ConnConfig.DefaultQueryExecMode = pgx.QueryExecModeSimpleProtocol - readConfig.ConnConfig.DefaultQueryExecMode = pgx.QueryExecModeSimpleProtocol - } + setDefaultQueryExecMode(writeConfig.ConnConfig) + setDefaultQueryExecMode(readConfig.ConnConfig) writeConfig.MinConns = int32(pg.maxIdleConnections) readConfig.MinConns = int32(pg.maxIdleConnections) @@ -75,14 +73,6 @@ func New(uri string, opts ...Option) (*Postgres, error) { writeConfig.MaxConnLifetimeJitter = time.Duration(0.2 * float64(pg.maxConnectionLifeTime)) readConfig.MaxConnLifetimeJitter = time.Duration(0.2 * float64(pg.maxConnectionLifeTime)) - if _, ok := readConfig.ConnConfig.Config.RuntimeParams["plan_cache_mode"]; !ok { - readConfig.ConnConfig.Config.RuntimeParams["plan_cache_mode"] = "force_custom_plan" - } - - if _, ok := writeConfig.ConnConfig.Config.RuntimeParams["plan_cache_mode"]; !ok { - writeConfig.ConnConfig.Config.RuntimeParams["plan_cache_mode"] = "force_custom_plan" - } - initialContext, cancelInit := context.WithTimeout(context.Background(), 5*time.Second) defer cancelInit() @@ -131,3 +121,27 @@ func (p *Postgres) IsReady(ctx context.Context) (bool, error) { } return true, nil } + +var queryExecModes = map[string]pgx.QueryExecMode{ + "cache_statement": pgx.QueryExecModeCacheStatement, + "cache_describe": pgx.QueryExecModeCacheDescribe, + "describe_exec": pgx.QueryExecModeDescribeExec, + "mode_exec": pgx.QueryExecModeExec, + "simple_protocol": pgx.QueryExecModeSimpleProtocol, +} + +func setDefaultQueryExecMode(config *pgx.ConnConfig) { + // Default mode if no specific mode is found in the connection string + defaultMode := "cache_statement" + + // Iterate through the map keys to check if any are mentioned in the connection string + for key := range queryExecModes { + if strings.Contains(config.ConnString(), "default_query_exec_mode="+key) { + config.DefaultQueryExecMode = queryExecModes[key] + return + } + } + + // Set to default mode if no matching mode is found + config.DefaultQueryExecMode = queryExecModes[defaultMode] +}