psql connection aliasing/management

.pgpass

# alias:northwind
db1.example.com:5432:northwind:postgres:hunter2

# alias:mainframe
db1.example.com:5432:mainframe:dave:password123

db1.example.com:5432:unaliased:postgres:

README.md

Setup:

• Add desired psql connections to ~/.pgpass with optional alias comments.
• Add the shell functions from functions.sh to your ~/.bashrc, ~/.zshrc, or wherever it can be sourced conveniently.
• Copy the pgpass.py script to ~/bin (or elsewhere, adjusting the value of PG_CONNECT_SCRIPT appropriately).
• Test usage as demonstrated in the CLI example section. You should be able to list and connect to your connections by alias using the pg-list and pg-connect functions.

CLI example

> source path/to/functions.sh
> pg-list
mainframe: dave:***@db1.example.com:5432/mainframe       
northwind: postgres:***@db1.example.com:5432/northwind   
unaliased: postgres:***@db1.example.com:5432/unaliased

> pg-connect northwind
pg-connect(northwind, psql -h db1.example.com -p 5432 -U postgres northwind)                                 
psql (9.4.4, server 8.0.2)                                                                            
SSL connection (protocol: TLSv1.2, cipher: ECDHE-RSA-AES256-GCM-SHA384, bits: 256, compression: on)   
Type "help" for help.                                                                                 

northwind=#                                                                                                

functions.sh

export PGPASSFILE=~/.pgpass                                                      
PG_CONNECT_SCRIPT=~/bin/pgpass.py                                       

function pg-connect () {                                                         
    PSQL_COMMAND=$($PG_CONNECT_SCRIPT print-connection-command $1)               
    echo "pg-connect($1, $PSQL_COMMAND)"                                         
    sh -c "$PSQL_COMMAND"                                                        
}                                                                                

function pg-list () {                                                            
    $PG_CONNECT_SCRIPT list                                                      
}                                                                                

pgpass.py

#!/Users/dave/projects/ci-etl/env/bin/python 
#!/usr/bin/python
import argparse
import collections
import os
import re
import sys


ConnectionLine = collections.namedtuple('ConnectionLine', ['host', 'port', 'database', 'user', 'password'])


def build_alias_map(filepath=None):
    '''
    Load a file with the format:
    |  # alias:{alias}
    |  host:port:database:user:password
    |  ...
    into a map like {'alias': ConnectionLine(host='host', port='port', database='database', user='user'), ...}
    '''
    filepath = filepath or os.environ.get('PGPASSFILE', '~/.pgpass')  
    alias_comment_pattern = re.compile(r'# alias:(?P<alias>[^\s]+)\s*')
    alias_map = {}
    unaliased_connections = []

    with open(filepath) as f:
        alias = None
        for line in f:
            m = alias_comment_pattern.match(line)
            if m:
                alias = m.group('alias')
                continue
            else:
                fields = line.strip().split(':')
                if len(fields) == 5:
                    connection_line = ConnectionLine(*fields)
                    if alias:
                        alias_map[alias] = connection_line
                    else:
                        unaliased_connections.append(connection_line)

            alias = None

    unaliased_map = {c.database: c for c in reversed(unaliased_connections)}

    merged_map = {}
    merged_map.update(unaliased_map)
    merged_map.update(alias_map)
    return merged_map


def get_connect_command(alias):
    alias_map = build_alias_map()
    connection_line = alias_map.get(alias)
    if not connection_line:
        raise KeyError("No connection found for alias: {}".format(alias))

    return 'psql -h {host} -p {port} -U {user} {database}'.format(
        user=connection_line.user,
        host=connection_line.host,
        port=connection_line.port,
        database=connection_line.database,
        )


def list_databases(filepath=None):
    alias_map = build_alias_map(filepath=filepath)
    aliases = sorted(alias_map.keys())
    for alias in aliases:
        connection_line = alias_map[alias]
        print '{alias}: {user}:***@{host}:{port}/{database}'.format(
            alias=alias,
            user=connection_line.user,
            host=connection_line.host,
            port=connection_line.port,
            database=connection_line.database,
            )


def dump_dbext():
    config_template = "let g:dbext_default_profile_{profile_name} = '{connection_string}'"
    connection_template = 'type={type}:user={user}:passwd={password}:dsnname={database}:host={host}:port={port}'

    alias_map = build_alias_map()
    aliases = sorted(alias_map.keys())
    for alias in aliases:
        connection_line = alias_map[alias]
        dbext_connection_string = connection_template.format(
            type='pgsql',
            user=connection_line.user,
            password=connection_line.password,
            host=connection_line.host,
            port=connection_line.port,
            database=connection_line.database,
            )
        config_line = config_template.format(
            profile_name=alias.lower().replace('-', '_'),
            connection_string=dbext_connection_string,
            )
        print config_line

        
def parse_opts(argv=None):
    argv = argv or sys.argv[1:]
    parser = argparse.ArgumentParser(description="Manage psql connections")
    subparsers = parser.add_subparsers()

    # pgpass.py list
    def cmd_list_databases(opts):
        list_databases()
    list_cmd = subparsers.add_parser('list')
    list_cmd.set_defaults(func=cmd_list_databases)

    # pgpass.py print-connection-command {alias}
    def cmd_print_connection_command(opts):
        print get_connect_command(opts.alias[0])
    connect_cmd = subparsers.add_parser('print-connection-command')
    connect_cmd.add_argument('alias', nargs=1)
    connect_cmd.set_defaults(func=cmd_print_connection_command)

    # pgpass.py dump-dbext
    def cmd_print_connection_command(opts):
        dump_dbext()
    connect_cmd = subparsers.add_parser('dump-dbext')
    connect_cmd.set_defaults(func=cmd_print_connection_command)

    return parser.parse_args(argv)


def main():
    opts = parse_opts()
    opts.func(opts)


if __name__ == '__main__':
    main()

is this web scale?

this is awesome, thanks!

This is really cool, thanks.
I made some minor modification to use 'pgcli' instead of 'psql' as it has more functionality if it is found on the system: https://gist.github.com/sirbijan/41a831f1aeee760e6e5f4afd4461fdb7/revisions

nice, works as intended! Protect me from screwin up prod when i think im in uat!

I like this, but it seems way overkill to use Python for something that could be done with a short bash script, so I made a version that just puts an alias before the standard .pgpass and just uses bash: https://gist.github.com/Fmstrat/ea6287a6d60e3e5f6c73e3bdd2f62331