Merge pull request #1239 from Permify/next

refactor(database): streamline default query execution mode setting
Permify · May 4, 2024 · a3f5546 · a3f5546
2 parents 7c2ecda + 42897a2
commit a3f5546
Show file tree

Hide file tree

Showing 9 changed files with 32 additions and 37 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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]
+}