Merge pull request #333 from dbcli/j-bennet/clean-pg-service-conf

Added "service" option.

3 files changed, 45 insertions, 12 deletions

diff --git a/pgcli/main.py b/pgcli/main.py
index 96c9f3df..e0bc2dce 100755
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -135,13 +135,17 @@ class PGCli(object):
         root_logger.debug('Initializing pgcli logging.')
         root_logger.debug('Log file %r.', log_file)
 
+    def connect_dsn(self, dsn):
+        self.connect(dsn=dsn)
+
     def connect_uri(self, uri):
         uri = urlparse(uri)
         database = uri.path[1:]  # ignore the leading fwd slash
         self.connect(database, uri.hostname, uri.username,
                      uri.port, uri.password)
 
-    def connect(self, database='', host='', user='', port='', passwd=''):
+    def connect(self, database='', host='', user='', port='', passwd='',
+                dsn=''):
         # Connect to the database.
 
         if not user:
@@ -174,13 +178,14 @@ class PGCli(object):
         # a password (no -w flag), prompt for a passwd and try again.
         try:
             try:
-                pgexecute = PGExecute(database, user, passwd, host, port)
+                pgexecute = PGExecute(database, user, passwd, host, port, dsn)
             except OperationalError as e:
                 if ('no password supplied' in utf8tounicode(e.args[0]) and
                         auto_passwd_prompt):
                     passwd = click.prompt('Password', hide_input=True,
                                           show_default=False, type=str)
-                    pgexecute = PGExecute(database, user, passwd, host, port)
+                    pgexecute = PGExecute(database, user, passwd, host, port,
+                                          dsn)
                 else:
                     raise e
 
@@ -455,6 +460,10 @@ def cli(database, user, host, port, prompt_passwd, never_prompt, dbname,
 
     if '://' in database:
         pgcli.connect_uri(database)
+    elif "=" in database:
+        pgcli.connect_dsn(database)
+    elif os.environ.get('PGSERVICE', None):
+        pgcli.connect_dsn('service={0}'.format(os.environ['PGSERVICE']))
     else:
         pgcli.connect(database, host, user, port)
 
diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py
index f857ba4d..2b226489 100644
--- a/pgcli/pgexecute.py
+++ b/pgcli/pgexecute.py
@@ -140,28 +140,42 @@ class PGExecute(object):
               AND n.nspname <> 'information_schema'
         ORDER BY 1, 2;'''
 
-    def __init__(self, database, user, password, host, port):
+    def __init__(self, database, user, password, host, port, dsn):
         self.dbname = database
         self.user = user
         self.password = password
         self.host = host
         self.port = port
+        self.dsn = dsn
         self.connect()
 
     def connect(self, database=None, user=None, password=None, host=None,
-            port=None):
+            port=None, dsn=None):
 
         db = (database or self.dbname)
         user = (user or self.user)
         password = (password or self.password)
         host = (host or self.host)
         port = (port or self.port)
-        conn = psycopg2.connect(
-                database=unicode2utf8(db),
-                user=unicode2utf8(user),
-                password=unicode2utf8(password),
-                host=unicode2utf8(host),
-                port=unicode2utf8(port))
+        dsn = (dsn or self.dsn)
+        if dsn:
+            if password:
+                dsn = "{0} password={1}".format(dsn, password)
+            conn = psycopg2.connect(dsn=unicode2utf8(dsn))
+            cursor = conn.cursor()
+            # When we connect using a DSN, we don't really know what db,
+            # user, etc. we connected to. Let's read it.
+            db = self._select_one(cursor, 'select current_database()')
+            user = self._select_one(cursor, 'select current_user')
+            host = self._select_one(cursor, 'select inet_server_addr()')
+            port = self._select_one(cursor, 'select inet_server_port()')
+        else:
+            conn = psycopg2.connect(
+                    database=unicode2utf8(db),
+                    user=unicode2utf8(user),
+                    password=unicode2utf8(password),
+                    host=unicode2utf8(host),
+                    port=unicode2utf8(port))
         conn.set_client_encoding('utf8')
         if hasattr(self, 'conn'):
             self.conn.close()
@@ -175,6 +189,16 @@ class PGExecute(object):
         register_json_typecasters(self.conn, self._json_typecaster)
         register_hstore_typecaster(self.conn)
 
+    def _select_one(self, cur, sql):
+        """
+        Helper method to run a select and retrieve a single field value
+        :param cur: cursor
+        :param sql: string
+        :return: string
+        """
+        cur.execute(sql)
+        return cur.fetchone()
+
     def _json_typecaster(self, json_data):
         """Interpret incoming JSON data as a string.
 
diff --git a/tests/conftest.py b/tests/conftest.py
index dfe86e8f..0cff635a 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -23,4 +23,4 @@ def cursor(connection):
 @pytest.fixture
 def executor(connection):
     return pgcli.pgexecute.PGExecute(database='_test_db', user=POSTGRES_USER,
-            host=POSTGRES_HOST, password=None, port=None)
+            host=POSTGRES_HOST, password=None, port=None, dsn=None)