Database Connection

Let’s create a connection string to the main database.

src_db_config = {
    'user': os.getenv('WWI_NEON_USER'),
    'pwd': os.getenv('WWI_NEON_PASSWORD'),
    'host': os.getenv('WWI_NEON_HOST'),
    'port': 5432, 
    'db': os.getenv('WWI_NEON_DB'),
    'sslmode': os.getenv('WWI_NEON_SSLMODE'),
    'channel_binding': os.getenv('WWI_NEON_CHANNEL_BINDING')
}

connection_string_wwi = (
    f"postgresql://{src_db_config['user']}:{src_db_config['pwd']}"
    f"@{src_db_config['host']}:{src_db_config['port']}/{src_db_config['db']}"
    f"?sslmode={src_db_config['sslmode']}"
    f"&channel_binding={src_db_config['channel_binding']}"
)

Let’s create a connection string to the analytical database. It is currently empty; we will load data into it later.

analytics_db_config = {
    'user': os.getenv('WWI_ANALYTICS_NEON_USER'),
    'pwd': os.getenv('WWI_ANALYTICS_NEON_PASSWORD'),
    'host': os.getenv('WWI_ANALYTICS_NEON_HOST'),
    'port': 5432, 
    'db': os.getenv('WWI_ANALYTICS_NEON_DB'),
    'sslmode': os.getenv('WWI_ANALYTICS_NEON_SSLMODE'),
    'channel_binding': os.getenv('WWI_ANALYTICS_NEON_CHANNEL_BINDING')
}

connection_string_wwi_analytics = (
    f"postgresql://{analytics_db_config['user']}:{analytics_db_config['pwd']}"
    f"@{analytics_db_config['host']}:{analytics_db_config['port']}/{analytics_db_config['db']}"
    f"?sslmode={analytics_db_config['sslmode']}"
    f"&channel_binding={analytics_db_config['channel_binding']}"
)

For the analytical database, we will need an engine; let’s create it now.

wwi_analytics_engine = create_engine(
    f"postgresql://{analytics_db_config['user']}:{analytics_db_config['pwd']}"
    f"@{analytics_db_config['host']}:{analytics_db_config['port']}/{analytics_db_config['db']}"
    , connect_args={
        'sslmode': os.getenv('WWI_ANALYTICS_NEON_SSLMODE'),
        'channel_binding': os.getenv('WWI_ANALYTICS_NEON_CHANNEL_BINDING')
    }
)

Let’s create a function to connect to the required database.

def con(db='src'):
    """
    Connects to the specified database using SQL magic.
    
    Parameters:
    -----------    
    db (str): Database to connect to - 'src' for source or 'dst' for destination.
                Defaults to 'src'.
    """
    if db == 'src':
        get_ipython().run_line_magic('sql', '$connection_string_wwi')
        print("Connected to srс")
    elif db == 'dst':
        get_ipython().run_line_magic('sql', '$connection_string_wwi_analytics')
        print("Connected to dst")
    else:
        raise ValueError(f"Unknown database identifier: '{db}'. Use 'src' or 'dst'")

Let’s create a function to close the connection.

def con_close(db='src'):
    """
    Closes the database connection for the specified database.
    
    Parameters:
    -----------
    db (str): Database to close - 'src' for source or 'dst' for destination.
                Defaults to 'src'.
    """
    # Get expected hosts from environment
    expected_hosts = {
        'src': os.getenv('WWI_NEON_HOST'),
        'dst': os.getenv('WWI_ANALYTICS_NEON_HOST')
    }
    
    if db not in expected_hosts:
        raise ValueError(f"Unknown db '{db}'. Use 'src' or 'dst'")
    
    # Get active connections
    connections = get_ipython().run_line_magic('sql', '--connections')
    
    # Find matching connection
    for url in connections.keys():
        if expected_hosts[db] in url:
            get_ipython().run_line_magic('sql', f'--close {url}')
            print(f"Connection closed for {db}")
            return
    
    print(f"No active connection found for {db}")

Let’s open connections to both databases for work.

con('src')

Connected to srс

con('dst')

Connected to dst