Merge pull request #40 from zalando/feature/db_connection_cmdline

Add -H and -P options to specify the host/port.

pg_view.py

@@ -5,7 +5,6 @@
 import re
 import sys
 import glob
-import getpass
 import logging
 from optparse import OptionParser
 from operator import itemgetter
@@ -104,7 +103,8 @@ def output_method_is_valid(method):
 def parse_args():
     '''parse command-line options'''
 
-    parser = OptionParser()
+    parser = OptionParser(add_help_option=False)
+    parser.add_option('-H', '--help', help='show_help', action='help')
     parser.add_option('-v', '--verbose', help='verbose mode', action='store_true', dest='verbose')
     parser.add_option('-i', '--instance', help='name of the instance to monitor', action='store', dest='instance')
     parser.add_option('-t', '--tick', help='tick length (in seconds)',
@@ -120,12 +120,16 @@ def parse_args():
                       dest='clear_screen')
     parser.add_option('-c', '--configuration-file', help='configuration file for PostgreSQL connections',
                       action='store', default='', dest='config_file')
-    parser.add_option('-p', '--pid', help='always track a given pid (may be used multiple times)',
+    parser.add_option('-P', '--pid', help='always track a given pid (may be used multiple times)',
                       action='append', type=int, default=[])
     parser.add_option('-U', '--username', help='database user name',
                       action='store', dest='username')
     parser.add_option('-d', '--dbname', help='database name to connect to',
                       action='store', dest='dbname')
+    parser.add_option('-h', '--host', help='database connection host '
+                      '(or a directory path for the unix socket connection)',
+                      action='store', dest='host')
+    parser.add_option('-p', '--port', help='database port number', action='store', dest='port')
 
     options, args = parser.parse_args()
     return options, args
@@ -2925,10 +2929,17 @@ def detect_db_port(socket_dir):
     return port
 
 
-def detect_default_user_database():
-    user = options.username or os.environ.get('PGUSER') or getpass.getuser()
-    database = options.dbname or os.environ.get('PGDATABASE') or user
-    return (user, database)
+def build_connection(host, port, user, database):
+    result = {}
+    if host:
+        result['host'] = host
+    if port:
+        result['port'] = port
+    if user:
+        result['user'] = user
+    if database:
+        result['database'] = database
+    return result
 
 
 def detect_postgres_version(work_directory):
@@ -3028,9 +3039,9 @@ def pick_connection_arguments(conn_args):
 
 def can_connect_with_connection_arguments(host, port):
     """ check that we can connect given the specified arguments """
-    user, database = detect_default_user_database()
+    conn = build_connection(host, port, options.username, options.dbname)
     try:
-        test_conn = psycopg2.connect('host={0} port={1} user={2} dbname={3}'.format(host, port, user, database))
+        test_conn = psycopg2.connect(**conn)
         test_conn.close()
     except psycopg2.OperationalError:
         return False
@@ -3073,30 +3084,17 @@ def detect_db_connection_arguments(work_directory, pid, version):
     return result
 
 
-def establish_user_defined_connection(dbname, args, clusters):
+def establish_user_defined_connection(instance, conn, clusters):
     """ connect the database and get all necessary options like pid and work_directory
         we use port, host and socket_directory, prefering socket over TCP connections
     """
 
     pgcon = None
-    # sanity check
-    if not (args.get('port') or args.get('host')):
-        missing = ('port' if not args.get('port') else 'host')
-        logger.error('Not all required connection arguments ' +
-                     '({0}) are specified for the database {1}, skipping it'.format(missing, dbname))
-        return None
-
-    port = args['port']
-    host = args['host']
-    user = args.get('user', 'postgres')
-    # the db is the actual database we connect to, while dbname is the name of the postgres cluster
-    db = args.get('dbname', 'postgres')
     # establish a new connection
     try:
-        pgcon = psycopg2.connect('host={0} port={1} user={2} dbname={3}'.format(host, port, user, db))
+        pgcon = psycopg2.connect(**conn)
     except Exception as e:
-        logger.error('failed to establish connection to {0} on port {1} user {2} database {3}'.format(host, port, user,
-                     dbname))
+        logger.error('failed to establish connection to {0} via {1}'.format(instance, conn))
         logger.error('PostgreSQL exception: {0}'.format(e))
         return None
     # get the database version from the pgcon properties
@@ -3109,23 +3107,23 @@ def establish_user_defined_connection(dbname, args, clusters):
     cur.close()
     pgcon.commit()
     # now, when we have the work directory, acquire the pid of the postmaster.
-    pid = read_postmaster_pid(work_directory, dbname)
+    pid = read_postmaster_pid(work_directory, instance)
     if pid is None:
-        logger.error('failed to read pid of the postmaster on {0}:{1}'.format(host, port))
+        logger.error('failed to read pid of the postmaster on {0}'.format(conn))
         return None
     # check that we don't have the same pid already in the accumulated results.
     # for instance, a user may specify 2 different set of connection options for
     # the same database (one for the unix_socket_directory and another for the host)
     pids = [opt['pid'] for opt in clusters if 'pid' in opt]
     if pid in pids:
-        duplicate_dbname = [opt['name'] for opt in clusters if 'pid' in opt and opt.get('pid', 0) == pid][0]
+        duplicate_instance = [opt['name'] for opt in clusters if 'pid' in opt and opt.get('pid', 0) == pid][0]
         logger.error('duplicate connection options detected ' +
-                     'for databases {0} and {1}, same pid {2}, skipping {0}'.format(dbname, duplicate_dbname, pid))
+                     'for databases {0} and {1}, same pid {2}, skipping {0}'.format(instance, duplicate_instance, pid))
         pgcon.close()
         return True
     # now we have all components to create the result
     result = {
-        'name': dbname,
+        'name': instance,
         'ver': dbver,
         'wd': work_directory,
         'pid': pid,
@@ -3287,15 +3285,27 @@ def main():
     clusters = []
 
     # now try to read the configuration file
-    config_data = (read_configuration(options.config_file) if options.config_file else None)
-    if config_data:
-        for dbname in config_data:
-            if user_dbname and dbname != user_dbname:
+    config = (read_configuration(options.config_file) if options.config_file else None)
+    if config:
+        for instance in config:
+            if user_dbname and instance != user_dbname:
                 continue
             # pass already aquired connections to make sure we only list unique clusters.
-            if not establish_user_defined_connection(dbname, config_data[dbname], clusters):
+            host = config[instance].get('host')
+            port = config[instance].get('port')
+            conn = build_connection(host, port,
+                                    config[instance].get('user'), config[instance].get('database'))
+
+            if not establish_user_defined_connection(instance, conn, clusters):
                 logger.error('failed to acquire details about ' +
-                             'the database cluster {0}, the server will be skipped'.format(dbname))
+                             'the database cluster {0}, the server will be skipped'.format(instance))
+    elif options.host:
+        port = options.port or "5432"
+        # try to connet to the database specified by command-line options
+        conn = build_connection(options.host, options.port, options.username, options.dbname)
+        instance = options.instance or "default"
+        if not establish_user_defined_connection(instance, conn, clusters):
+            logger.error("unable to continue with cluster {0}".format(instance))
     else:
         # do autodetection
         postmasters = get_postmasters_directories()
@@ -3315,8 +3325,8 @@ def main():
                     continue
                 host = conndata['host']
                 port = conndata['port']
-                user, database = detect_default_user_database()
-                pgcon = psycopg2.connect('host={0} port={1} user={2} dbname={3}'.format(host, port, user, database))
+                conn = build_connection(host, port, options.username, options.dbname)
+                pgcon = psycopg2.connect(**conn)
             except Exception as e:
                 logger.error('PostgreSQL exception {0}'.format(e))
                 pgcon = None