diff --git a/persistent/internal/driver/postgres/basic_operations.go b/persistent/internal/driver/postgres/basic_operations.go index 0bae3be1..cb4f02ee 100644 --- a/persistent/internal/driver/postgres/basic_operations.go +++ b/persistent/internal/driver/postgres/basic_operations.go @@ -15,7 +15,7 @@ import ( // Insert adds one or more objects into the database in a single batch operation. // Returns an error if the input is empty or the insert fails. func (d *driver) Insert(ctx context.Context, objects ...model.DBObject) error { - if d.db == nil { + if d.writeDB == nil { return errors.New(types.ErrorSessionClosed) } if len(objects) == 0 { @@ -52,7 +52,7 @@ func (d *driver) Insert(ctx context.Context, objects ...model.DBObject) error { } tableName := objs[0].TableName() - if err := d.db.WithContext(ctx).Table(tableName).Create(sliceValue.Interface()).Error; err != nil { + if err := d.writeDB.WithContext(ctx).Table(tableName).Create(sliceValue.Interface()).Error; err != nil { return err } } @@ -73,7 +73,7 @@ func (d *driver) Delete(ctx context.Context, object model.DBObject, filters ...m } // Start building the query with the table name - db := d.db.WithContext(ctx).Table(tableName) + db := d.writeDB.WithContext(ctx).Table(tableName) // If we have a filter, use our translator function if len(filters) == 1 { db, err = d.translateQuery(db, filters[0], object) @@ -111,7 +111,7 @@ func (d *driver) Update(ctx context.Context, object model.DBObject, filters ...m return errors.New(types.ErrorMultipleDBM) } - tx := d.db.WithContext(ctx).Table(tableName) + tx := d.writeDB.WithContext(ctx).Table(tableName) // Apply filters if len(filters) == 1 { @@ -149,7 +149,7 @@ for updating a collection of specific records with different values. */ func (d *driver) BulkUpdate(ctx context.Context, objects []model.DBObject, filters ...model.DBM) error { // Basic validation - if d.db == nil { + if d.writeDB == nil { return errors.New(types.ErrorSessionClosed) } if len(objects) == 0 { @@ -160,7 +160,7 @@ func (d *driver) BulkUpdate(ctx context.Context, objects []model.DBObject, filte } // Start a transaction - tx := d.db.WithContext(ctx).Begin() + tx := d.writeDB.WithContext(ctx).Begin() if tx.Error != nil { return tx.Error } @@ -272,7 +272,7 @@ func (d *driver) UpdateAll(ctx context.Context, row model.DBObject, query, updat } // Start a transaction - tx := d.db.WithContext(ctx).Begin() + tx := d.writeDB.WithContext(ctx).Begin() if tx.Error != nil { return tx.Error } @@ -284,7 +284,7 @@ func (d *driver) UpdateAll(ctx context.Context, row model.DBObject, query, updat panic(r) // re-throw panic after rollback } }() - db := d.db.WithContext(ctx).Table(tableName) + db := d.writeDB.WithContext(ctx).Table(tableName) // Check if query is empty hasFilter := false @@ -343,7 +343,7 @@ func (d *driver) Upsert(ctx context.Context, row model.DBObject, query, update m return err } - tx := d.db.WithContext(ctx).Begin() + tx := d.writeDB.WithContext(ctx).Begin() if tx.Error != nil { return tx.Error } @@ -430,43 +430,6 @@ func (d *driver) fetchUpdatedRow(tx *gorm.DB, table string, query model.DBM, row return db.First(row).Error } -func ensureID(originalID model.ObjectID, row model.DBObject, query model.DBM) { - if originalID != "" { - row.SetObjectID(originalID) - } else if idVal, ok := query["id"].(string); ok && idVal != "" { - row.SetObjectID(model.ObjectIDHex(idVal)) - } - if row.GetObjectID() == "" { - row.SetObjectID(model.NewObjectID()) - } -} - -func cloneDBObject(row model.DBObject) model.DBObject { - newRow := reflect.New(reflect.TypeOf(row).Elem()).Interface().(model.DBObject) - newRow.SetObjectID(row.GetObjectID()) - return newRow -} - -func mergeQueryFields(row model.DBObject, query model.DBM) { - for k, v := range query { - if strings.HasPrefix(k, "_") || k == "$or" { - continue - } - setField(row, k, v) // keeps reflection logic isolated - } -} - -func (d *driver) ensureID(originalID model.ObjectID, row model.DBObject, query model.DBM) { - if originalID != "" { - row.SetObjectID(originalID) - } else if qid, ok := query["id"]; ok { - if sid, ok2 := qid.(string); ok2 && sid != "" { - row.SetObjectID(model.ObjectIDHex(sid)) - - } - } -} - // Helper function to set a field in a struct using reflection func setField(obj interface{}, name string, value interface{}) { structValue := reflect.ValueOf(obj) diff --git a/persistent/internal/driver/postgres/basic_operations_test.go b/persistent/internal/driver/postgres/basic_operations_test.go index 68f44a91..47f014cd 100644 --- a/persistent/internal/driver/postgres/basic_operations_test.go +++ b/persistent/internal/driver/postgres/basic_operations_test.go @@ -549,73 +549,3 @@ func TestUpsert(t *testing.T) { assert.Equal(t, 50, resultItem.Value) }) } - -func TestCloneDBObject(t *testing.T) { - original := &TestObject{ - Name: "Original", - Value: 42, - CreatedAt: time.Now(), - } - original.SetObjectID(model.NewObjectID()) - - clone := cloneDBObject(original) - - // Ensure it's a different pointer - assert.NotSame(t, original, clone) - - // Ensure it has the same ID - assert.Equal(t, original.GetObjectID(), clone.GetObjectID()) - - // Ensure other fields are zeroed (because cloneDBObject only copies ID) - cloneObj, ok := clone.(*TestObject) - require.True(t, ok) - - assert.Equal(t, "", cloneObj.Name) - assert.Equal(t, 0, cloneObj.Value) - assert.WithinDuration(t, time.Time{}, cloneObj.CreatedAt, time.Second) -} - -func TestMergeQueryFields(t *testing.T) { - obj := &TestObject{ - Name: "Initial", - Value: 10, - CreatedAt: time.Now(), - } - - query := model.DBM{ - "name": "Updated Name", - "value": 42, - "_limit": 100, // should be ignored - "$or": []model.DBM{}, // should be ignored - "extra_field": "Extra", // will only work if TestObject has this field; otherwise ignored - } - - mergeQueryFields(obj, query) - - // Check that allowed fields were updated - assert.Equal(t, "Updated Name", obj.Name) - assert.Equal(t, 42, obj.Value) - -} - -func TestEnsureID(t *testing.T) { - driver, _ := setupTest(t) - defer teardownTest(t, driver) - - // Case 1: originalID is provided → should preserve it - obj1 := &TestObject{} - origID := model.NewObjectID() - driver.ensureID(origID, obj1, model.DBM{}) - assert.Equal(t, origID, obj1.GetObjectID()) - - // Case 2: originalID is empty, but query contains "id" - obj2 := &TestObject{} - queryID := model.NewObjectID() - driver.ensureID("", obj2, model.DBM{"id": queryID.Hex()}) - assert.Equal(t, queryID, obj2.GetObjectID()) - - // Case 3: neither originalID nor query["id"] → ID should remain empty - obj3 := &TestObject{} - driver.ensureID("", obj3, model.DBM{}) - assert.Equal(t, model.ObjectID(""), obj3.GetObjectID()) -} diff --git a/persistent/internal/driver/postgres/driver.go b/persistent/internal/driver/postgres/driver.go index d20695f9..1741310f 100644 --- a/persistent/internal/driver/postgres/driver.go +++ b/persistent/internal/driver/postgres/driver.go @@ -31,7 +31,7 @@ func NewPostgresDriver(opts *types.ClientOpts) (*driver, error) { func (d *driver) validateDBAndTable(object model.DBObject) (string, error) { // Check if the database connection is valid - if d.db == nil { + if d.writeDB == nil || d.readDB == nil { return "", errors.New(types.ErrorSessionClosed) } diff --git a/persistent/internal/driver/postgres/driver_test.go b/persistent/internal/driver/postgres/driver_test.go index 6f123e0c..6685811c 100644 --- a/persistent/internal/driver/postgres/driver_test.go +++ b/persistent/internal/driver/postgres/driver_test.go @@ -122,7 +122,7 @@ func TestValidateDBAndTable(t *testing.T) { // Create a driver with a valid connection driver, _ := setupTest(t) - // Close the connection to simulate a nil db + // Close the connection to simulate a nil writeDB driver.Close() // Create a mock object with a valid table name diff --git a/persistent/internal/driver/postgres/indexes.go b/persistent/internal/driver/postgres/indexes.go index e01cdebf..ebea7c4a 100644 --- a/persistent/internal/driver/postgres/indexes.go +++ b/persistent/internal/driver/postgres/indexes.go @@ -11,6 +11,8 @@ import ( "strings" ) +var sanitizerRegex = regexp.MustCompile(`^[a-zA-Z_][a-zA-Z0-9_]*$`) + type IndexRow struct { IndexName string ColumnName string @@ -20,13 +22,6 @@ type IndexRow struct { Comment *string // Using pointer for nullable string } -func sanitizeIdentifier(s string) (string, error) { - if matched := regexp.MustCompile(`^[a-zA-Z_][a-zA-Z0-9_]*$`).MatchString(s); !matched { - return "", fmt.Errorf("invalid identifier: %s", s) - } - return pq.QuoteIdentifier(s), nil // use pq or pgx quoting -} - // CreateIndex creates a database index on the specified table for the given fields. // Returns an error if index creation fails. func (d *driver) CreateIndex(ctx context.Context, row model.DBObject, index model.Index) error { @@ -135,7 +130,7 @@ func (d *driver) CreateIndex(ctx context.Context, row model.DBObject, index mode createSQL = strings.Replace(createSQL, "CREATE INDEX", "CREATE INDEX CONCURRENTLY", 1) } - if err := d.db.WithContext(ctx).Exec(createSQL).Error; err != nil { + if err := d.writeDB.WithContext(ctx).Exec(createSQL).Error; err != nil { return fmt.Errorf("failed to create index: %w", err) } @@ -150,7 +145,7 @@ func (d *driver) CreateIndex(ctx context.Context, row model.DBObject, index mode PRIMARY KEY (table_name, index_name) ) ` - if err := d.db.WithContext(ctx).Exec(metadataSQL).Error; err != nil { + if err := d.writeDB.WithContext(ctx).Exec(metadataSQL).Error; err != nil { return fmt.Errorf("failed to create metadata table: %w", err) } @@ -160,7 +155,7 @@ func (d *driver) CreateIndex(ctx context.Context, row model.DBObject, index mode ON CONFLICT (table_name, index_name) DO UPDATE SET is_ttl = TRUE, ttl_seconds = EXCLUDED.ttl_seconds ` - if err := d.db.WithContext(ctx).Exec(ttlSQL, tableName, indexName, index.TTL).Error; err != nil { + if err := d.writeDB.WithContext(ctx).Exec(ttlSQL, tableName, indexName, index.TTL).Error; err != nil { return fmt.Errorf("failed to store TTL metadata: %w", err) } } @@ -232,7 +227,7 @@ func (d *driver) GetIndexes(ctx context.Context, row model.DBObject) ([]model.In ` var rows []IndexRow - err = d.db.WithContext(ctx).Raw(query, tableName).Scan(&rows).Error + err = d.readDB.WithContext(ctx).Raw(query, tableName).Scan(&rows).Error if err != nil { return nil, fmt.Errorf("failed to query indexes: %w", err) } @@ -274,7 +269,7 @@ func (d *driver) GetIndexes(ctx context.Context, row model.DBObject) ([]model.In func (d *driver) tableExists(ctx context.Context, tableName string) (bool, error) { var exists bool - err := d.db.WithContext(ctx).Raw("SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = ?)", tableName).Scan(&exists).Error + err := d.readDB.WithContext(ctx).Raw("SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = ?)", tableName).Scan(&exists).Error if err != nil { return false, fmt.Errorf("failed to check if table exists: %w", err) } @@ -316,7 +311,7 @@ func (d *driver) CleanIndexes(ctx context.Context, row model.DBObject) error { var indexNames []string // Execute the query - err = d.db.WithContext(ctx).Raw(query, tableName).Scan(&indexNames).Error + err = d.writeDB.WithContext(ctx).Raw(query, tableName).Scan(&indexNames).Error if err != nil { return fmt.Errorf("failed to query indexes: %w", err) } @@ -326,7 +321,7 @@ func (d *driver) CleanIndexes(ctx context.Context, row model.DBObject) error { } // Start a transaction - tx := d.db.WithContext(ctx).Begin() + tx := d.writeDB.WithContext(ctx).Begin() if tx.Error != nil { return fmt.Errorf("failed to begin transaction: %w", err) } @@ -365,10 +360,17 @@ func (d *driver) indexExists(ctx context.Context, tableName, indexName string) ( ` var exists bool - err := d.db.WithContext(ctx).Raw(query, tableName, indexName).Scan(&exists).Error + err := d.readDB.WithContext(ctx).Raw(query, tableName, indexName).Scan(&exists).Error if err != nil { return false, err } return exists, nil } + +func sanitizeIdentifier(s string) (string, error) { + if matched := sanitizerRegex.MatchString(s); !matched { + return "", fmt.Errorf("invalid identifier: %s", s) + } + return pq.QuoteIdentifier(s), nil // use pq or pgx quoting +} diff --git a/persistent/internal/driver/postgres/indexes_test.go b/persistent/internal/driver/postgres/indexes_test.go index c37575ad..b57071bc 100644 --- a/persistent/internal/driver/postgres/indexes_test.go +++ b/persistent/internal/driver/postgres/indexes_test.go @@ -18,7 +18,7 @@ func TestCreateIndex(t *testing.T) { // Helper function to clean up test data cleanupTestData := func(tableName string) { - err := driver.db.WithContext(ctx).Exec(fmt.Sprintf("DROP TABLE IF EXISTS %s", tableName)).Error + err := driver.writeDB.WithContext(ctx).Exec(fmt.Sprintf("DROP TABLE IF EXISTS %s", tableName)).Error if err != nil { t.Logf("Error cleaning up test data: %v", err) } diff --git a/persistent/internal/driver/postgres/lifecycle.go b/persistent/internal/driver/postgres/lifecycle.go index 0161e0fe..005b99e7 100644 --- a/persistent/internal/driver/postgres/lifecycle.go +++ b/persistent/internal/driver/postgres/lifecycle.go @@ -15,58 +15,113 @@ import ( ) type lifeCycle struct { - db *gorm.DB - connectionString string - dbName string - sqlDB *sql.DB -} + writeDB *gorm.DB + readDB *gorm.DB -// Connect initializes a new database connection using the provided client options. -// Returns an error if the connection cannot be established. -func (l *lifeCycle) Connect(opts *types.ClientOpts) error { - if opts == nil { - return errors.New("nil opts") - } + dbName string + writeSQLDB *sql.DB + readSQLDB *sql.DB +} - // Use the connection string exactly as provided - dsn := strings.TrimSpace(opts.ConnectionString) - if dsn == "" { - return errors.New("empty connection string") - } +func (l *lifeCycle) establishConnection(dsn string, opts *types.ClientOpts) (*gorm.DB, *sql.DB, error) { + // Create dialect from DSN + dialect := postgres.New(postgres.Config{ + DSN: dsn, + }) - // Open GORM with PostgreSQL driver - db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{}) + // Open connection + db, err := gorm.Open(dialect, &gorm.Config{}) if err != nil { - return fmt.Errorf("gorm open: %w", err) + return nil, nil, fmt.Errorf("failed to open connection: %w", err) } - l.db = db - // Get underlying sql.DB + // Get underlying SQL DB sqlDB, err := db.DB() if err != nil { - l.db = nil - return fmt.Errorf("getting sql.DB: %w", err) + return nil, nil, fmt.Errorf("failed to get sql.DB: %w", err) } - l.sqlDB = sqlDB - // Set connection lifetime if provided + // Configure connection if opts.ConnectionTimeout > 0 { - l.sqlDB.SetConnMaxLifetime(time.Duration(opts.ConnectionTimeout) * time.Second) + sqlDB.SetConnMaxLifetime(time.Duration(opts.ConnectionTimeout) * time.Second) } - // Ping to verify connection + // Ping database to verify connection pingTimeout := 5 * time.Second if opts.ConnectionTimeout > 0 { pingTimeout = time.Duration(opts.ConnectionTimeout) * time.Second } + ctx, cancel := context.WithTimeout(context.Background(), pingTimeout) defer cancel() - if err := l.sqlDB.PingContext(ctx); err != nil { - l.db = nil - return fmt.Errorf("ping db: %w", err) + if err := sqlDB.PingContext(ctx); err != nil { + err := sqlDB.Close() + if err != nil { + return nil, nil, err + } + return nil, nil, fmt.Errorf("database ping failed: %w", err) + } + + return db, sqlDB, nil +} + +// Connect initializes a new database connection using the provided client options. +// Returns an error if the connection cannot be established. +func (l *lifeCycle) Connect(opts *types.ClientOpts) error { + if opts == nil { + return errors.New("nil opts") + } + + // Validate and normalize connection strings + writeDSN, readDSN, err := validateAndNormalizeDSNs(opts.ConnectionString, opts.ReadConnectionString) + if err != nil { + return err } - // Extract database name (supports both DSN and URL formats) + // Establish write connection + l.writeDB, l.writeSQLDB, err = l.establishConnection(writeDSN, opts) + if err != nil { + return fmt.Errorf("failed to establish write connection: %w", err) + } + + // Extract database name from the write connection + l.extractDBName(writeDSN) + + // Handle read connection + if readDSN == writeDSN { + // Use write connection for reads if they're the same + l.readDB = l.writeDB + l.readSQLDB = l.writeSQLDB + } else { + // Establish separate read connection + l.readDB, l.readSQLDB, err = l.establishConnection(readDSN, opts) + if err != nil { + // Clean up write connection + err := l.writeSQLDB.Close() + if err != nil { + return errors.New("failed to close write database") + } + l.writeDB = nil + l.writeSQLDB = nil + return errors.New("failed to establish read connection") + } + } + return nil +} + +func (l *lifeCycle) closeWriteConnection() error { + if l.writeSQLDB != nil { + err := l.writeSQLDB.Close() + if err != nil { + return errors.New("failed to close write connection") + } + l.writeDB = nil + l.writeSQLDB = nil + } + return nil +} + +func (l *lifeCycle) extractDBName(dsn string) { if parts := strings.Fields(dsn); len(parts) > 0 { for _, part := range parts { if strings.HasPrefix(part, "dbname=") { @@ -80,19 +135,39 @@ func (l *lifeCycle) Connect(opts *types.ClientOpts) error { l.dbName = strings.TrimPrefix(u.Path, "/") } } - - return nil } // Close terminates the active database connection. // Returns an error if the connection cannot be closed properly. func (l *lifeCycle) Close() error { - err := l.sqlDB.Close() - if err != nil { - return err + var err error + connectionClosed := false + + // Close write connection + // Close write connection + if l.writeSQLDB != nil { + err = l.writeSQLDB.Close() + connectionClosed = true + l.writeDB = nil + l.writeSQLDB = nil } - l.db = nil - return nil + + // Only close read connection if it's different from write connection + // or if we haven't closed any connection yet + if l.readSQLDB != nil && (!connectionClosed || l.readSQLDB != l.writeSQLDB) { + readErr := l.readSQLDB.Close() + // If we didn't have an error from closing write connection, + // use the error from closing read connection + if err == nil { + err = readErr + } + } + + // Always nil the read pointers + l.readDB = nil + l.readSQLDB = nil + + return err } // DBType returns the type of database managed by this lifecycle. @@ -104,8 +179,8 @@ func (l *lifeCycle) DBType() utils.DBType { // Ping checks the health of the database connection. // Returns an error if the database is unreachable or not responding. func (d *driver) Ping(ctx context.Context) error { - // Check if the database connection is valid - if d.db == nil { + // Check if connections exist + if d.writeDB == nil || d.readDB == nil { return errors.New(types.ErrorSessionClosed) } @@ -113,16 +188,29 @@ func (d *driver) Ping(ctx context.Context) error { return errors.New(types.ErrorNilContext) } - // Get the underlying *sql.DB from GORM - sqlDB, err := d.db.DB() + // Ping write connection + writeSQLDB, err := d.writeDB.DB() if err != nil { - return fmt.Errorf("failed to get database connection: %w", err) + return fmt.Errorf("failed to get write database connection: %w", err) + } + + if err := writeSQLDB.PingContext(ctx); err != nil { + return fmt.Errorf("write database ping failed: %w", err) + } + + // If read and write are the same connection, we're done + if d.readDB == d.writeDB { + return nil } - // Use the standard library's PingContext method - err = sqlDB.PingContext(ctx) + // Ping read connection + readSQLDB, err := d.readDB.DB() if err != nil { - return fmt.Errorf("failed to ping database: %w", err) + return fmt.Errorf("failed to get read database connection: %w", err) + } + + if err := readSQLDB.PingContext(ctx); err != nil { + return fmt.Errorf("read database ping failed: %w", err) } return nil @@ -132,7 +220,7 @@ func (d *driver) Ping(ctx context.Context) error { // Returns an error if the operation fails. func (d *driver) DropDatabase(_ context.Context) error { // Check if the database connection is valid - if d.db == nil { + if d.writeDB == nil { return errors.New(types.ErrorSessionClosed) } @@ -147,8 +235,8 @@ func (d *driver) DropDatabase(_ context.Context) error { return fmt.Errorf("failed to close connection: %w", err) } - // Set the driver's db to nil to prevent further use - d.db = nil + // Set the driver's writeDB to nil to prevent further use + d.writeDB = nil // Return a special error with instructions return fmt.Errorf( @@ -158,4 +246,20 @@ func (d *driver) DropDatabase(_ context.Context) error { dbNameToDelete, dbNameToDelete) } +// validateAndNormalizeDSNs validates connection strings and applies fallback logic. +// It returns the normalized write and read DSNs, or an error if validation fails. +func validateAndNormalizeDSNs(writeConnStr, readConnStr string) (string, string, error) { + // Validate write connection string + if writeConnStr == "" { + return "", "", errors.New("write connection string is required") + } + + // If no separate read connection specified, use the write connection for reads + if readConnStr == "" { + readConnStr = writeConnStr + } + + return writeConnStr, readConnStr, nil +} + var _ types.StorageLifecycle = &lifeCycle{} diff --git a/persistent/internal/driver/postgres/lifecycle_test.go b/persistent/internal/driver/postgres/lifecycle_test.go index de7bdf74..be5ec81d 100644 --- a/persistent/internal/driver/postgres/lifecycle_test.go +++ b/persistent/internal/driver/postgres/lifecycle_test.go @@ -92,7 +92,7 @@ func TestPing(t *testing.T) { // Create a driver instance driver, ctx := setupTest(t) - // Explicitly close the session to make d.db nil + // Explicitly close the session to make d.writeDB nil err := driver.Close() assert.NoError(t, err, "Failed to close the driver session") @@ -137,16 +137,16 @@ func TestLifeCycleConnect(t *testing.T) { assert.NoError(t, err, "Connect should not return an error with valid options") // Verify that the connection is established - assert.NotNil(t, lc.db, "Database connection should not be nil after successful connection") + assert.NotNil(t, lc.writeDB, "Database connection should not be nil after successful connection") // Verify that the connection works by pinging the database ctx := context.Background() - err = lc.db.WithContext(ctx).Exec("SELECT 1").Error + err = lc.writeDB.WithContext(ctx).Exec("SELECT 1").Error assert.NoError(t, err, "Should be able to execute a simple query after connection") // Clean up - if lc.db != nil { - sqlDB, err := lc.db.DB() + if lc.writeDB != nil { + sqlDB, err := lc.writeDB.DB() if err == nil { sqlDB.Close() } @@ -174,7 +174,7 @@ func TestLifeCycleConnect(t *testing.T) { assert.Error(t, err, "Connect should return an error with invalid connection string") // Verify that the database connection is nil - assert.Nil(t, lc.db, "Database connection should be nil after failed connection") + assert.Nil(t, lc.writeDB, "Database connection should be nil after failed connection") }) // Test case 3: Failed connection due to nil options @@ -189,7 +189,7 @@ func TestLifeCycleConnect(t *testing.T) { assert.Error(t, err, "Connect should return an error with nil options") // Verify that the database connection is nil - assert.Nil(t, lc.db, "Database connection should be nil after failed connection") + assert.Nil(t, lc.writeDB, "Database connection should be nil after failed connection") }) // Test case 4: Failed connection due to empty connection string @@ -210,7 +210,7 @@ func TestLifeCycleConnect(t *testing.T) { assert.Error(t, err, "Connect should return an error with empty connection string") // Verify that the database connection is nil - assert.Nil(t, lc.db, "Database connection should be nil after failed connection") + assert.Nil(t, lc.writeDB, "Database connection should be nil after failed connection") }) } @@ -226,7 +226,47 @@ func TestDropDatabase(t *testing.T) { driver, _ := setupTest(t) err := driver.DropDatabase(context.Background()) - assert.Nil(t, driver.db) + assert.Nil(t, driver.writeDB) assert.Error(t, err) }) } + +func TestValidateAndNormalizeDSNs(t *testing.T) { + // Test case 1: Both write and read DSNs provided + writeDSN := "host=primary.db.example.com dbname=testdb user=writer" + readDSN := "host=replica.db.example.com dbname=testdb user=reader" + + w, r, err := validateAndNormalizeDSNs(writeDSN, readDSN) + assert.NoError(t, err) + assert.Equal(t, writeDSN, w) + assert.Equal(t, readDSN, r) + + // Test case 2: Only write DSN provided + writeDSN = "host=primary.db.example.com dbname=testdb user=writer" + readDSN = "" + + w, r, err = validateAndNormalizeDSNs(writeDSN, readDSN) + assert.NoError(t, err) + assert.Equal(t, writeDSN, w) + assert.Equal(t, writeDSN, r, "Read DSN should fall back to write DSN when not provided") + + // Test case 3: No write DSN provided (error case) + writeDSN = "" + readDSN = "host=replica.db.example.com dbname=testdb user=reader" + + w, r, err = validateAndNormalizeDSNs(writeDSN, readDSN) + assert.Error(t, err) + assert.Equal(t, "", w) + assert.Equal(t, "", r) + assert.Contains(t, err.Error(), "write connection string is required") + + // Test case 4: Both DSNs empty (error case) + writeDSN = "" + readDSN = "" + + w, r, err = validateAndNormalizeDSNs(writeDSN, readDSN) + assert.Error(t, err) + assert.Equal(t, "", w) + assert.Equal(t, "", r) + assert.Contains(t, err.Error(), "write connection string is required") +} diff --git a/persistent/internal/driver/postgres/query.go b/persistent/internal/driver/postgres/query.go index 9ddacec0..1a31b879 100644 --- a/persistent/internal/driver/postgres/query.go +++ b/persistent/internal/driver/postgres/query.go @@ -28,7 +28,7 @@ func (d *driver) Query(ctx context.Context, object model.DBObject, result interf return errors.New("result must be a pointer") } - db := d.db.WithContext(ctx).Table(tableName) + db := d.readDB.WithContext(ctx).Table(tableName) db, err = d.translateQuery(db, filter, object) if err != nil { return err @@ -76,7 +76,7 @@ func (d *driver) Count(ctx context.Context, row model.DBObject, filters ...model return 0, errors.New(types.ErrorCollectionNotFound) } - db := d.db.WithContext(ctx).Table(tableName) + db := d.readDB.WithContext(ctx).Table(tableName) // If we have a filter, use our translator function if len(filters) == 1 { // Add _count flag to the filter to ensure proper handling in translateQuery @@ -121,7 +121,7 @@ func (d *driver) Aggregate(ctx context.Context, row model.DBObject, pipeline []m } // Execute the query using GORM - rows, err := d.db.WithContext(ctx).Raw(sqlQuery, args...).Rows() + rows, err := d.readDB.WithContext(ctx).Raw(sqlQuery, args...).Rows() if err != nil { return nil, fmt.Errorf("failed to execute aggregation query: %w", err) } @@ -400,7 +400,7 @@ func (d *driver) translateQuery(db *gorm.DB, q model.DBM, result interface{}) (* AND tablename LIKE ? ` - rows, err := d.db.Raw(query, tablePattern).Rows() + rows, err := d.readDB.Raw(query, tablePattern).Rows() if err != nil { return nil, fmt.Errorf("failed to get sharded tables: %w", err) } diff --git a/persistent/internal/driver/postgres/query_test.go b/persistent/internal/driver/postgres/query_test.go index 00af9544..8f08275a 100644 --- a/persistent/internal/driver/postgres/query_test.go +++ b/persistent/internal/driver/postgres/query_test.go @@ -20,7 +20,7 @@ func TestQuery(t *testing.T) { // Helper function to clean up test data cleanupTestData := func(tableName string) { - err := driver.db.WithContext(ctx).Exec(fmt.Sprintf("DELETE FROM %s", tableName)).Error + err := driver.writeDB.WithContext(ctx).Exec(fmt.Sprintf("DELETE FROM %s", tableName)).Error if err != nil { t.Logf("Error cleaning up test data: %v", err) } @@ -170,7 +170,7 @@ func TestCount(t *testing.T) { // Helper function to clean up test data cleanupTestData := func(tableName string) { - err := driver.db.WithContext(ctx).Exec(fmt.Sprintf("DELETE FROM %s", tableName)).Error + err := driver.writeDB.WithContext(ctx).Exec(fmt.Sprintf("DELETE FROM %s", tableName)).Error if err != nil { t.Logf("Error cleaning up test data: %v", err) } @@ -325,7 +325,7 @@ func TestAggregate(t *testing.T) { // Helper function to clean up test data cleanupTestData := func(tableName string) { - err := driver.db.WithContext(ctx).Exec(fmt.Sprintf("DELETE FROM %s", tableName)).Error + err := driver.writeDB.WithContext(ctx).Exec(fmt.Sprintf("DELETE FROM %s", tableName)).Error if err != nil { t.Logf("Error cleaning up test data: %v", err) } @@ -929,7 +929,7 @@ func TestApplyMongoUpdateOperators(t *testing.T) { // Helper function to clean up test data cleanupTestData := func(tableName string) { - err := driver.db.WithContext(ctx).Exec(fmt.Sprintf("DELETE FROM %s", tableName)).Error + err := driver.writeDB.WithContext(ctx).Exec(fmt.Sprintf("DELETE FROM %s", tableName)).Error if err != nil { t.Logf("Error cleaning up test data: %v", err) } @@ -978,7 +978,7 @@ func TestApplyMongoUpdateOperators(t *testing.T) { } // Apply the update operators - db := driver.db.WithContext(ctx).Table(tableName) + db := driver.writeDB.WithContext(ctx).Table(tableName) updatedDB, updates, err := driver.applyMongoUpdateOperators(db, update) require.NoError(t, err, "applyMongoUpdateOperators should not return an error") @@ -992,7 +992,7 @@ func TestApplyMongoUpdateOperators(t *testing.T) { // Verify the object was updated var updatedObj TestObject - err = driver.db.WithContext(ctx).Table(tableName).Where("id = ?", id).First(&updatedObj).Error + err = driver.readDB.WithContext(ctx).Table(tableName).Where("id = ?", id).First(&updatedObj).Error require.NoError(t, err, "Failed to retrieve updated object") assert.Equal(t, "Updated Name", updatedObj.Name, "Name should be updated") @@ -1014,7 +1014,7 @@ func TestApplyMongoUpdateOperators(t *testing.T) { } // Apply the update operators - db := driver.db.WithContext(ctx).Table(tableName) + db := driver.writeDB.WithContext(ctx).Table(tableName) updatedDB, updates, err := driver.applyMongoUpdateOperators(db, update) require.NoError(t, err, "applyMongoUpdateOperators should not return an error") @@ -1024,7 +1024,7 @@ func TestApplyMongoUpdateOperators(t *testing.T) { // Verify the object was updated var updatedObj TestObject - err = driver.db.WithContext(ctx).Table(tableName).Where("id = ?", id).First(&updatedObj).Error + err = driver.readDB.WithContext(ctx).Table(tableName).Where("id = ?", id).First(&updatedObj).Error require.NoError(t, err, "Failed to retrieve updated object") assert.Equal(t, 15, updatedObj.Value, "Value should be incremented by 5") @@ -1047,7 +1047,7 @@ func TestApplyMongoUpdateOperators(t *testing.T) { } // Apply the update operators - db := driver.db.WithContext(ctx).Table(tableName) + db := driver.writeDB.WithContext(ctx).Table(tableName) updatedDB, updates, err := driver.applyMongoUpdateOperators(db, update) require.NoError(t, err, "applyMongoUpdateOperators should not return an error") @@ -1057,7 +1057,7 @@ func TestApplyMongoUpdateOperators(t *testing.T) { // Verify the object was updated var updatedObj TestObject - err = driver.db.WithContext(ctx).Table(tableName).Where("id = ?", id).First(&updatedObj).Error + err = driver.readDB.WithContext(ctx).Table(tableName).Where("id = ?", id).First(&updatedObj).Error require.NoError(t, err, "Failed to retrieve updated object") assert.Equal(t, "Updated Multiple", updatedObj.Name, "Name should be updated") @@ -1069,7 +1069,7 @@ func TestApplyMongoUpdateOperators(t *testing.T) { update := model.DBM{} // Apply the update operators - db := driver.db.WithContext(ctx).Table("test_table") + db := driver.writeDB.WithContext(ctx).Table("test_table") updatedDB, updates, err := driver.applyMongoUpdateOperators(db, update) require.NoError(t, err, "applyMongoUpdateOperators should not return an error with empty update") @@ -1094,7 +1094,7 @@ func TestApplyMongoUpdateOperators(t *testing.T) { } // Apply the update operators - db := driver.db.WithContext(ctx).Table(tableName) + db := driver.writeDB.WithContext(ctx).Table(tableName) updatedDB, updates, err := driver.applyMongoUpdateOperators(db, update) require.NoError(t, err, "applyMongoUpdateOperators should not return an error") @@ -1104,7 +1104,7 @@ func TestApplyMongoUpdateOperators(t *testing.T) { // Verify the object was updated var updatedObj TestObject - err = driver.db.WithContext(ctx).Table(tableName).Where("id = ?", id).First(&updatedObj).Error + err = driver.readDB.WithContext(ctx).Table(tableName).Where("id = ?", id).First(&updatedObj).Error require.NoError(t, err, "Failed to retrieve updated object") assert.Equal(t, 30, updatedObj.Value, "Value should be multiplied by 3") @@ -1124,7 +1124,7 @@ func TestApplyMongoUpdateOperators(t *testing.T) { } // Apply the update operators - db := driver.db.WithContext(ctx).Table(tableName) + db := driver.writeDB.WithContext(ctx).Table(tableName) updatedDB, updates, err := driver.applyMongoUpdateOperators(db, update) require.NoError(t, err, "applyMongoUpdateOperators should not return an error") @@ -1134,7 +1134,7 @@ func TestApplyMongoUpdateOperators(t *testing.T) { // Verify the object was updated var updatedObj TestObject - err = driver.db.WithContext(ctx).Table(tableName).Where("id = ?", id).First(&updatedObj).Error + err = driver.readDB.WithContext(ctx).Table(tableName).Where("id = ?", id).First(&updatedObj).Error require.NoError(t, err, "Failed to retrieve updated object") assert.Equal(t, 3, updatedObj.Value, "Value should be 3") @@ -1154,7 +1154,7 @@ func TestApplyMongoUpdateOperators(t *testing.T) { } // Apply the update operators - db := driver.db.WithContext(ctx).Table(tableName) + db := driver.writeDB.WithContext(ctx).Table(tableName) updatedDB, updates, err := driver.applyMongoUpdateOperators(db, update) require.NoError(t, err, "applyMongoUpdateOperators should not return an error") @@ -1164,7 +1164,7 @@ func TestApplyMongoUpdateOperators(t *testing.T) { // Verify the object was updated var updatedObj TestObject - err = driver.db.WithContext(ctx).Table(tableName).Where("id = ?", id).First(&updatedObj).Error + err = driver.readDB.WithContext(ctx).Table(tableName).Where("id = ?", id).First(&updatedObj).Error require.NoError(t, err, "Failed to retrieve updated object") assert.Equal(t, 300, updatedObj.Value, "Value should be 300") @@ -1186,7 +1186,7 @@ func TestApplyMongoUpdateOperators(t *testing.T) { } // Apply the update operators - db := driver.db.WithContext(ctx).Table(tableName) + db := driver.writeDB.WithContext(ctx).Table(tableName) updatedDB, updates, err := driver.applyMongoUpdateOperators(db, update) require.NoError(t, err, "applyMongoUpdateOperators should not return an error") @@ -1196,7 +1196,7 @@ func TestApplyMongoUpdateOperators(t *testing.T) { // Verify the object was updated var updatedObj TestObject - err = driver.db.WithContext(ctx).Table(tableName).Where("id = ?", id).First(&updatedObj).Error + err = driver.readDB.WithContext(ctx).Table(tableName).Where("id = ?", id).First(&updatedObj).Error require.NoError(t, err, "Failed to retrieve updated object") assert.Equal(t, "", updatedObj.Name, "Name should be empty") @@ -1218,7 +1218,7 @@ func TestApplyMongoUpdateOperators(t *testing.T) { } // Apply the update operators - db := driver.db.WithContext(ctx).Table(tableName) + db := driver.writeDB.WithContext(ctx).Table(tableName) updatedDB, updates, err := driver.applyMongoUpdateOperators(db, update) require.NoError(t, err, "applyMongoUpdateOperators should not return an error") @@ -1228,7 +1228,7 @@ func TestApplyMongoUpdateOperators(t *testing.T) { // Verify the object was updated var updatedObj TestObject - err = driver.db.WithContext(ctx).Table(tableName).Where("id = ?", id).First(&updatedObj).Error + err = driver.readDB.WithContext(ctx).Table(tableName).Where("id = ?", id).First(&updatedObj).Error require.NoError(t, err, "Failed to retrieve updated object") assert.True(t, updatedObj.CreatedAt.After(pastTime), "CreatedAt should be updated to a newer time") }) @@ -1242,7 +1242,7 @@ func TestApplyMongoUpdateOperators(t *testing.T) { } // Apply the update operators - db := driver.db.WithContext(ctx).Table("test_table") + db := driver.writeDB.WithContext(ctx).Table("test_table") _, _, err := driver.applyMongoUpdateOperators(db, update) // Check if the function returns an error for unsupported operator @@ -1259,11 +1259,11 @@ func TestTranslateQuery(t *testing.T) { // Ensure the test table exists tableName := "test_objects" - err := driver.db.WithContext(ctx).AutoMigrate(&TestObject{}) + err := driver.writeDB.WithContext(ctx).AutoMigrate(&TestObject{}) require.NoError(t, err, "Failed to create test table") // Clean up any existing data - err = driver.db.WithContext(ctx).Exec("DELETE FROM " + tableName).Error + err = driver.writeDB.WithContext(ctx).Exec("DELETE FROM " + tableName).Error require.NoError(t, err, "Failed to clean up test table") // Create test data @@ -1297,7 +1297,7 @@ func TestTranslateQuery(t *testing.T) { // Insert test data for _, obj := range testData { - err := driver.db.WithContext(ctx).Table(tableName).Create(&obj).Error + err := driver.writeDB.WithContext(ctx).Table(tableName).Create(&obj).Error require.NoError(t, err, "Failed to insert test data") } @@ -1464,7 +1464,7 @@ func TestTranslateQuery(t *testing.T) { for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { // Create a base DB query - db := driver.db.WithContext(ctx).Table(tableName) + db := driver.readDB.WithContext(ctx).Table(tableName) // Apply the translateQuery function translatedDB, err := driver.translateQuery(db, tc.query, &TestObject{}) @@ -1482,7 +1482,7 @@ func TestTranslateQuery(t *testing.T) { // Test count flag separately t.Run("Count", func(t *testing.T) { // Create a base DB query - db := driver.db.WithContext(ctx).Table(tableName) + db := driver.readDB.WithContext(ctx).Table(tableName) // Apply the translateQuery function with count flag query := model.DBM{ @@ -1524,7 +1524,7 @@ func TestTranslateQueryWithShardingEnabled(t *testing.T) { shardTableName := baseTableName + "_" + date.Format("20060102") // Create the sharded table - err := driver.db.WithContext(ctx).Exec(fmt.Sprintf(` + err := driver.writeDB.WithContext(ctx).Exec(fmt.Sprintf(` CREATE TABLE IF NOT EXISTS %s ( id TEXT PRIMARY KEY, name TEXT, @@ -1538,7 +1538,7 @@ func TestTranslateQueryWithShardingEnabled(t *testing.T) { // Insert test data into the sharded table for j := 1; j <= 3; j++ { id := model.NewObjectID() - err := driver.db.WithContext(ctx).Exec(fmt.Sprintf(` + err := driver.writeDB.WithContext(ctx).Exec(fmt.Sprintf(` INSERT INTO %s (id, name, value, category, created_at) VALUES (?, ?, ?, ?, ?) `, shardTableName), id.Hex(), fmt.Sprintf("Shard %d-%d", i, j), j*10, @@ -1558,7 +1558,7 @@ func TestTranslateQueryWithShardingEnabled(t *testing.T) { } // Create a base DB query - db := driver.db.WithContext(ctx).Table(baseTableName) + db := driver.readDB.WithContext(ctx).Table(baseTableName) // Apply the translateQuery function testObj := &TestObject{} @@ -1585,7 +1585,7 @@ func TestTranslateQueryWithShardingEnabled(t *testing.T) { for i := 0; i <= 3; i++ { date := startDate.Add(time.Duration(i*24) * time.Hour) shardTableName := baseTableName + "_" + date.Format("20060102") - err := driver.db.WithContext(ctx).Exec(fmt.Sprintf("DROP TABLE IF EXISTS %s", shardTableName)).Error + err := driver.writeDB.WithContext(ctx).Exec(fmt.Sprintf("DROP TABLE IF EXISTS %s", shardTableName)).Error require.NoError(t, err, "Failed to drop sharded table") } } diff --git a/persistent/internal/driver/postgres/schema.go b/persistent/internal/driver/postgres/schema.go index 4a348b9d..1fa074a4 100644 --- a/persistent/internal/driver/postgres/schema.go +++ b/persistent/internal/driver/postgres/schema.go @@ -66,8 +66,9 @@ type BasicInfo struct { // HasTable checks if a table with the given name exists in the database. // Returns true if the table exists, otherwise false with any error encountered. func (d *driver) HasTable(ctx context.Context, tableName string) (bool, error) { + db := d.readDB // Check if the database connection is valid - if d.db == nil { + if db == nil { return false, errors.New(types.ErrorSessionClosed) } @@ -76,12 +77,13 @@ func (d *driver) HasTable(ctx context.Context, tableName string) (bool, error) { return false, errors.New(types.ErrorEmptyTableName) } - return d.db.Migrator().HasTable(tableName), nil + return db.Migrator().HasTable(tableName), nil } // DBTableStats returns statistics for the given table. // Provides details like row count, size, and other metadata as a map. func (d *driver) DBTableStats(ctx context.Context, row model.DBObject) (model.DBM, error) { + db := d.readDB tableName, err := d.validateDBAndTable(row) if err != nil { return nil, err @@ -121,7 +123,7 @@ func (d *driver) DBTableStats(ctx context.Context, row model.DBObject) (model.DB ` var basicStats BasicStats - err = d.db.WithContext(ctx).Raw(basicStatsQuery, tableName).Scan(&basicStats).Error + err = db.WithContext(ctx).Raw(basicStatsQuery, tableName).Scan(&basicStats).Error if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return nil, fmt.Errorf("table %s not found", tableName) @@ -182,7 +184,7 @@ func (d *driver) DBTableStats(ctx context.Context, row model.DBObject) (model.DB ` var indexStatsRows []IndexStats - err = d.db.WithContext(ctx).Raw(indexStatsQuery, tableName).Scan(&indexStatsRows).Error + err = db.WithContext(ctx).Raw(indexStatsQuery, tableName).Scan(&indexStatsRows).Error if err != nil { return nil, fmt.Errorf("failed to get index statistics: %w", err) } @@ -231,7 +233,7 @@ func (d *driver) DBTableStats(ctx context.Context, row model.DBObject) (model.DB ` var columnStatsRows []ColumnStats - err = d.db.WithContext(ctx).Raw(columnStatsQuery, tableName, tableName).Scan(&columnStatsRows).Error + err = db.WithContext(ctx).Raw(columnStatsQuery, tableName, tableName).Scan(&columnStatsRows).Error if err != nil { return nil, fmt.Errorf("failed to get column statistics: %w", err) } @@ -285,12 +287,13 @@ func (d *driver) DBTableStats(ctx context.Context, row model.DBObject) (model.DB // GetTables retrieves the list of all table names in the current database. // Returns a slice of table names or an error if the query fails. func (d *driver) GetTables(ctx context.Context) ([]string, error) { - if d.db == nil { + db := d.readDB + if db == nil { return []string{}, errors.New(types.ErrorSessionClosed) } var tables []string - err := d.db.Raw(`SELECT tablename FROM pg_tables WHERE schemaname = 'public'`).Scan(&tables).Error + err := db.Raw(`SELECT tablename FROM pg_tables WHERE schemaname = 'public'`).Scan(&tables).Error if err != nil { return nil, fmt.Errorf("failed to get tables: %w", err) } @@ -301,7 +304,9 @@ func (d *driver) GetTables(ctx context.Context) ([]string, error) { // DropTable removes a table by its name. // Returns the number of tables dropped and any error encountered. func (d *driver) DropTable(ctx context.Context, name string) (int, error) { - if d.db == nil { + writeDB := d.writeDB + readDB := d.readDB + if writeDB == nil || readDB == nil { return 0, errors.New(types.ErrorSessionClosed) } @@ -319,12 +324,12 @@ func (d *driver) DropTable(ctx context.Context, name string) (int, error) { // This is to return the number of affected rows var rowCount int64 countQuery := fmt.Sprintf("SELECT COUNT(*) FROM %s", name) - err = d.db.WithContext(ctx).Raw(countQuery).Scan(&rowCount).Error + err = readDB.WithContext(ctx).Raw(countQuery).Scan(&rowCount).Error if err != nil { return 0, err } - err = d.db.Migrator().DropTable(ctx, name) + err = writeDB.Migrator().DropTable(ctx, name) if err != nil { return 0, fmt.Errorf("failed to drop table %s: %w", name, err) } @@ -334,13 +339,14 @@ func (d *driver) DropTable(ctx context.Context, name string) (int, error) { // Drop removes the table associated with the given object. // Permanently deletes its schema and all stored data. func (d *driver) Drop(ctx context.Context, object model.DBObject) error { + db := d.writeDB // Check if the database connection is valid tableName, err := d.validateDBAndTable(object) if err != nil { return err } - err = d.db.WithContext(ctx).Migrator().DropTable(tableName) + err = db.WithContext(ctx).Migrator().DropTable(tableName) if err != nil { return fmt.Errorf("failed to drop table %s: %w", tableName, err) } @@ -352,7 +358,7 @@ func (d *driver) Drop(ctx context.Context, object model.DBObject) error { // Ensures tables and indexes are created or updated based on the models. func (d *driver) Migrate(ctx context.Context, objects []model.DBObject, options ...model.DBM) error { // Check if the database connection is valid - if d.db == nil { + if d.writeDB == nil { return errors.New(types.ErrorSessionClosed) } @@ -367,7 +373,7 @@ func (d *driver) Migrate(ctx context.Context, objects []model.DBObject, options } // Use GORM's context - db := d.db.WithContext(ctx) + writeDB := d.writeDB.WithContext(ctx) // Process each object for _, obj := range objects { // Get the table name @@ -380,12 +386,12 @@ func (d *driver) Migrate(ctx context.Context, objects []model.DBObject, options } // Check if the table already exists - tableExists := db.Migrator().HasTable(tableName) + tableExists := writeDB.Migrator().HasTable(tableName) if tableExists { continue // Skip if table already exists } - err := db.Table(tableName).AutoMigrate(obj) + err := writeDB.Table(tableName).AutoMigrate(obj) if err != nil { return fmt.Errorf("failed to migrate table %s: %w", tableName, err) } @@ -421,7 +427,8 @@ func (d *driver) Migrate(ctx context.Context, objects []model.DBObject, options // GetDatabaseInfo returns metadata about the connected database. // Provides details such as type, version, and connection information. func (d *driver) GetDatabaseInfo(ctx context.Context) (utils.Info, error) { - if d.db == nil { + readDB := d.readDB + if readDB == nil { return utils.Info{}, errors.New(types.ErrorSessionClosed) } @@ -441,7 +448,7 @@ func (d *driver) GetDatabaseInfo(ctx context.Context) (utils.Info, error) { current_setting('max_connections') AS max_connections ` var basicInfo BasicInfo - err := d.db.WithContext(ctx).Raw(basicInfoQuery).Scan(&basicInfo).Error + err := readDB.WithContext(ctx).Raw(basicInfoQuery).Scan(&basicInfo).Error if err != nil { return utils.Info{}, fmt.Errorf("failed to get database info: %w", err) } @@ -470,7 +477,7 @@ func (d *driver) GetDatabaseInfo(ctx context.Context) (utils.Info, error) { ` var connectionCount int - err = d.db.WithContext(ctx).Raw(connectionCountQuery).Scan(&connectionCount).Error + err = readDB.WithContext(ctx).Raw(connectionCountQuery).Scan(&connectionCount).Error if err == nil { info.CurrentConnections = connectionCount } @@ -487,7 +494,7 @@ func (d *driver) GetDatabaseInfo(ctx context.Context) (utils.Info, error) { ` var tableCount int - err = d.db.WithContext(ctx).Raw(tableCountQuery).Scan(&tableCount).Error + err = readDB.WithContext(ctx).Raw(tableCountQuery).Scan(&tableCount).Error if err == nil { info.TableCount = tableCount } diff --git a/persistent/internal/driver/postgres/schema_test.go b/persistent/internal/driver/postgres/schema_test.go index 3dd46b3b..25159827 100644 --- a/persistent/internal/driver/postgres/schema_test.go +++ b/persistent/internal/driver/postgres/schema_test.go @@ -82,7 +82,7 @@ func TestHasTable(t *testing.T) { // If not, you may need to modify or remove this test // Create a schema if it doesn't exist - err := driver.db.WithContext(ctx).Exec("CREATE SCHEMA IF NOT EXISTS test_schema").Error + err := driver.writeDB.WithContext(ctx).Exec("CREATE SCHEMA IF NOT EXISTS test_schema").Error require.NoError(t, err) schemaItem := &TestObject{TableNameValue: "test_schema.schema_table"} @@ -106,7 +106,7 @@ func TestHasTable(t *testing.T) { assert.NoError(t, err) // Drop the schema - err = driver.db.WithContext(ctx).Exec("DROP SCHEMA IF EXISTS test_schema CASCADE").Error + err = driver.writeDB.WithContext(ctx).Exec("DROP SCHEMA IF EXISTS test_schema CASCADE").Error assert.NoError(t, err) }) } @@ -190,7 +190,8 @@ func TestGetTables(t *testing.T) { }) t.Run("nil connection", func(t *testing.T) { - driver.db = nil + driver.writeDB = nil + driver.readDB = nil tables, err := driver.GetTables(ctx) assert.Error(t, err) assert.Len(t, tables, 0) @@ -261,7 +262,7 @@ func TestDropTable(t *testing.T) { // Test case 4: Drop a table in a different schema t.Run("DropTableInDifferentSchema", func(t *testing.T) { // Create a test schema - err := driver.db.WithContext(ctx).Exec("CREATE SCHEMA IF NOT EXISTS test_schema").Error + err := driver.writeDB.WithContext(ctx).Exec("CREATE SCHEMA IF NOT EXISTS test_schema").Error require.NoError(t, err, "Failed to create test schema") // Create a test table in the schema @@ -284,7 +285,7 @@ func TestDropTable(t *testing.T) { assert.False(t, exists, "Table in schema should not exist after being dropped") // Drop the schema - err = driver.db.WithContext(ctx).Exec("DROP SCHEMA IF EXISTS test_schema CASCADE").Error + err = driver.writeDB.WithContext(ctx).Exec("DROP SCHEMA IF EXISTS test_schema CASCADE").Error assert.NoError(t, err, "Failed to drop test schema") }) @@ -351,7 +352,8 @@ func TestDropTable(t *testing.T) { }) t.Run("Nil connection", func(t *testing.T) { - driver.db = nil + driver.writeDB = nil + driver.readDB = nil rows, err := driver.DropTable(ctx, "test_drop_table") assert.Error(t, err) assert.Equal(t, 0, rows) @@ -463,7 +465,7 @@ func TestDrop(t *testing.T) { // Test case 7: Drop a table in a different schema t.Run("DropTableInDifferentSchema", func(t *testing.T) { // Create a test schema - err := driver.db.WithContext(ctx).Exec("CREATE SCHEMA IF NOT EXISTS test_schema").Error + err := driver.writeDB.WithContext(ctx).Exec("CREATE SCHEMA IF NOT EXISTS test_schema").Error require.NoError(t, err, "Failed to create test schema") // Create a test table in the schema @@ -486,7 +488,7 @@ func TestDrop(t *testing.T) { assert.False(t, exists, "Table in schema should not exist after being dropped") // Drop the schema - err = driver.db.WithContext(ctx).Exec("DROP SCHEMA IF EXISTS test_schema CASCADE").Error + err = driver.writeDB.WithContext(ctx).Exec("DROP SCHEMA IF EXISTS test_schema CASCADE").Error assert.NoError(t, err, "Failed to drop test schema") }) } @@ -641,7 +643,8 @@ func TestMigrate(t *testing.T) { t.Run("Nil connection", func(t *testing.T) { testObj := &TestObject{TableNameValue: "test_migrate"} defer cleanupTables([]model.DBObject{testObj}) - driver.db = nil + driver.writeDB = nil + driver.readDB = nil // Migrate the table err := driver.Migrate(ctx, []model.DBObject{testObj}) diff --git a/persistent/internal/driver/postgres/utils.go b/persistent/internal/driver/postgres/utils.go index 428fe692..39f1b60f 100644 --- a/persistent/internal/driver/postgres/utils.go +++ b/persistent/internal/driver/postgres/utils.go @@ -2,7 +2,9 @@ package postgres import ( "encoding/json" + "github.com/TykTechnologies/storage/persistent/model" "reflect" + "strings" ) // Helper functions @@ -59,3 +61,33 @@ func getCollectionName(result interface{}) (string, bool) { return "", false } + +func ensureID(originalID model.ObjectID, row model.DBObject, query model.DBM) { + if originalID != "" { + row.SetObjectID(originalID) + } else if qid, ok := query["id"]; ok { + if sid, ok2 := qid.(string); ok2 && sid != "" { + row.SetObjectID(model.ObjectIDHex(sid)) + } + } + + // Always ensure there's an ID + if row.GetObjectID() == "" { + row.SetObjectID(model.NewObjectID()) + } +} + +func cloneDBObject(row model.DBObject) model.DBObject { + newRow := reflect.New(reflect.TypeOf(row).Elem()).Interface().(model.DBObject) + newRow.SetObjectID(row.GetObjectID()) + return newRow +} + +func mergeQueryFields(row model.DBObject, query model.DBM) { + for k, v := range query { + if strings.HasPrefix(k, "_") || k == "$or" { + continue + } + setField(row, k, v) // keeps reflection logic isolated + } +} diff --git a/persistent/internal/driver/postgres/utils_test.go b/persistent/internal/driver/postgres/utils_test.go index 518514b6..0977f43f 100644 --- a/persistent/internal/driver/postgres/utils_test.go +++ b/persistent/internal/driver/postgres/utils_test.go @@ -4,7 +4,9 @@ package postgres import ( + "github.com/TykTechnologies/storage/persistent/model" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "testing" "time" ) @@ -288,3 +290,74 @@ func TestGetCollectionName(t *testing.T) { } }) } + +func TestCloneDBObject(t *testing.T) { + original := &TestObject{ + Name: "Original", + Value: 42, + CreatedAt: time.Now(), + } + original.SetObjectID(model.NewObjectID()) + + clone := cloneDBObject(original) + + // Ensure it's a different pointer + assert.NotSame(t, original, clone) + + // Ensure it has the same ID + assert.Equal(t, original.GetObjectID(), clone.GetObjectID()) + + // Ensure other fields are zeroed (because cloneDBObject only copies ID) + cloneObj, ok := clone.(*TestObject) + require.True(t, ok) + + assert.Equal(t, "", cloneObj.Name) + assert.Equal(t, 0, cloneObj.Value) + assert.WithinDuration(t, time.Time{}, cloneObj.CreatedAt, time.Second) +} + +func TestMergeQueryFields(t *testing.T) { + obj := &TestObject{ + Name: "Initial", + Value: 10, + CreatedAt: time.Now(), + } + + query := model.DBM{ + "name": "Updated Name", + "value": 42, + "_limit": 100, // should be ignored + "$or": []model.DBM{}, // should be ignored + "extra_field": "Extra", // will only work if TestObject has this field; otherwise ignored + } + + mergeQueryFields(obj, query) + + // Check that allowed fields were updated + assert.Equal(t, "Updated Name", obj.Name) + assert.Equal(t, 42, obj.Value) + +} + +func TestEnsureID(t *testing.T) { + driver, _ := setupTest(t) + defer teardownTest(t, driver) + + // Case 1: originalID is provided → should preserve it + obj1 := &TestObject{} + origID := model.NewObjectID() + ensureID(origID, obj1, model.DBM{}) + ensureID(origID, obj1, model.DBM{}) + assert.Equal(t, origID, obj1.GetObjectID()) + + // Case 2: originalID is empty, but the query contains "id" + obj2 := &TestObject{} + queryID := model.NewObjectID() + ensureID("", obj2, model.DBM{"id": queryID.Hex()}) + assert.Equal(t, queryID, obj2.GetObjectID()) + + // Case 3: neither originalID nor query["id"] → a new ID is generated + obj3 := &TestObject{} + ensureID("", obj3, model.DBM{}) + assert.NotEqual(t, "", obj3.GetObjectID()) +} diff --git a/persistent/internal/types/client_options.go b/persistent/internal/types/client_options.go index b827fb6b..7d98458d 100644 --- a/persistent/internal/types/client_options.go +++ b/persistent/internal/types/client_options.go @@ -18,6 +18,25 @@ type ClientOpts struct { // ConnectionString is the expression used to connect to a storage db server. // It contains parameters such as username, hostname, password and port ConnectionString string + // ReadConnectionString specifies the connection string for read operations. + // This enables read/write separation, allowing read queries to be directed to + // read replicas while write operations use the primary database specified by ConnectionString. + // + // If ReadConnectionString is empty, the system will use ConnectionString for both + // read and write operations, effectively using a single database connection. + // + // Read/write separation is particularly useful for: + // - Scaling read-heavy workloads by distributing reads across replicas + // - Reducing load on the primary database + // - Improving read performance by connecting to geographically closer replicas + // - Implementing high availability patterns + // + // Example usage: + // opts := &ClientOpts{ + // ConnectionString: "host=primary.db.example.com dbname=mydb user=writer", + // ReadConnectionString: "host=replica.db.example.com dbname=mydb user=reader", + // } + ReadConnectionString string // UseSSL is SSL connection is required to connect UseSSL bool // This setting allows the use of self-signed certificates when connecting to an encrypted storage database.