PostgreSQL Connection Management

/**
 * PostgreSQL connection implementation
 */
public class PostgreSQLConnection extends Connection<PostgreSQLDatabase> {
    /**
     * Creates a new PostgreSQL connection wrapper
     * @param database The PostgreSQL database instance
     * @param connection The raw JDBC connection
     */
    protected PostgreSQLConnection(PostgreSQLDatabase database, java.sql.Connection connection);
    
    /**
     * Checks if this connection is running on AWS RDS
     * @return true if running on AWS RDS
     */
    public boolean isAwsRds();
    
    /**
     * Restores the connection to its original state
     * Resets the role to the original user role
     * @throws SQLException if role cannot be restored
     */
    protected void doRestoreOriginalState() throws SQLException;
    
    /**
     * Gets the current schema for this connection
     * @return PostgreSQLSchema instance for current schema
     * @throws SQLException if current schema cannot be determined
     */
    public Schema doGetCurrentSchema() throws SQLException;
    
    /**
     * Gets the current schema name or search path
     * @return Current search path setting
     * @throws SQLException if search path cannot be retrieved
     */
    protected String getCurrentSchemaNameOrSearchPath() throws SQLException;
    
    /**
     * Changes the current schema by modifying the search path
     * @param schema The schema to change to
     */
    public void changeCurrentSchemaTo(Schema schema);
    
    /**
     * Sets the search path for the connection
     * @param schema Comma-separated schema names for search path
     * @throws SQLException if search path cannot be set
     */
    public void doChangeCurrentSchemaOrSearchPathTo(String schema) throws SQLException;
    
    /**
     * Creates a PostgreSQL schema instance
     * @param name Schema name
     * @return PostgreSQLSchema instance
     */
    public Schema getSchema(String name);
    
    /**
     * Executes a callable with advisory table locking
     * Uses PostgreSQL advisory locks to prevent concurrent access
     * @param table The table to lock
     * @param callable The operation to execute while locked
     * @param <T> Return type of the callable
     * @return Result from the callable execution
     */
    public <T> T lock(Table table, Callable<T> callable);
}

import org.flywaydb.database.postgresql.PostgreSQLConnection;
import java.util.concurrent.Callable;

// Connection is typically obtained from PostgreSQLDatabase
PostgreSQLConnection connection = (PostgreSQLConnection) database.getConnection();

// Check AWS RDS status
boolean isRds = connection.isAwsRds();

// Get current schema
Schema currentSchema = connection.doGetCurrentSchema();
System.out.println("Current schema: " + currentSchema.getName());

// Change to a different schema
Schema targetSchema = connection.getSchema("public");
connection.changeCurrentSchemaTo(targetSchema);

// Execute with advisory locking
Table migrationTable = database.getTable("flyway_schema_history");
String result = connection.lock(migrationTable, new Callable<String>() {
    @Override
    public String call() throws Exception {
        // Perform migration operations with table locked
        return "Migration completed";
    }
});

/**
 * PostgreSQL advisory lock template for safe concurrent operations
 */
public class PostgreSQLAdvisoryLockTemplate {
    /**
     * Creates an advisory lock template
     * @param configuration Flyway configuration
     * @param jdbcTemplate JDBC template for database operations
     * @param lockId Unique lock identifier (typically table name hash)
     */
    public PostgreSQLAdvisoryLockTemplate(Configuration configuration, 
                                         JdbcTemplate jdbcTemplate, 
                                         int lockId);
    
    /**
     * Executes a callable while holding the advisory lock
     * @param callable Operation to execute
     * @param <T> Return type
     * @return Result from callable execution
     */
    public <T> T execute(Callable<T> callable);
}

import org.flywaydb.database.postgresql.PostgreSQLAdvisoryLockTemplate;

// Create lock template for a specific table
int lockId = "my_table".hashCode();
PostgreSQLAdvisoryLockTemplate lockTemplate = 
    new PostgreSQLAdvisoryLockTemplate(configuration, jdbcTemplate, lockId);

// Execute operation with locking
String result = lockTemplate.execute(() -> {
    // Database operations that need exclusive access
    jdbcTemplate.execute("INSERT INTO my_table VALUES (1, 'data')");
    return "Operation completed";
});

// Get current schema with error handling
try {
    Schema current = connection.doGetCurrentSchema();
    if (current.exists()) {
        System.out.println("Working with schema: " + current.getName());
    }
} catch (SQLException e) {
    // Handle schema detection errors
}

// Switch schemas safely
Schema newSchema = connection.getSchema("new_schema");
if (newSchema.exists()) {
    connection.changeCurrentSchemaTo(newSchema);
} else {
    newSchema.create();
    connection.changeCurrentSchemaTo(newSchema);
}

// Get current search path
String searchPath = connection.getCurrentSchemaNameOrSearchPath();
System.out.println("Current search path: " + searchPath);

// Set custom search path
connection.doChangeCurrentSchemaOrSearchPathTo("app_schema,public");

// Restore original search path (handled automatically on connection close)

PostgreSQLConnection connection = (PostgreSQLConnection) database.getConnection();

if (connection.isAwsRds()) {
    // RDS-specific handling
    System.out.println("Running on AWS RDS - using optimized settings");
    // Certain operations may behave differently on RDS
}

// Original role is automatically stored during connection creation
// Role changes during migration are automatically restored when connection closes

// Manual role restoration (typically not needed)
connection.doRestoreOriginalState();  // Resets to original role

try {
    Schema schema = connection.doGetCurrentSchema();
} catch (FlywayException e) {
    // Thrown when search path is empty and current schema cannot be determined
    System.err.println("Schema detection failed: " + e.getMessage());
    // Consider setting explicit schema in configuration
}

try {
    connection.doChangeCurrentSchemaOrSearchPathTo("nonexistent_schema");
} catch (FlywaySqlException e) {
    // Thrown when schema change fails
    System.err.println("Schema change failed: " + e.getMessage());
}