≈Merge branch 'master' of github.com:dbcli/pgcli into 816

4 files changed, 27 insertions, 6 deletions

diff --git a/AUTHORS b/AUTHORS
index b5878064..990f0bb5 100644
--- a/AUTHORS
+++ b/AUTHORS
@@ -69,6 +69,7 @@ Contributors:
     * Andrew Speed
     * Dmitry B
     * Isank
+    * Marcin Sztolcman
 
 Creator:
 --------
diff --git a/changelog.rst b/changelog.rst
index 4f7dbd2f..87eb9aa2 100644
--- a/changelog.rst
+++ b/changelog.rst
@@ -1,5 +1,6 @@
 Upcoming
 ========
+* Use other prompt (prompt_dsn) when connecting using --dsn parameter. (Thanks: `Marcin Sztolcman`_)
 
 TODO
 
@@ -763,3 +764,4 @@ Improvements:
 .. _`Alexander Schmolck`: https://github.com/aschmolck
 .. _`Andrew Speed`: https://github.com/AndrewSpeed
 .. _`Dmitry B`: https://github.com/oxitnik
+.. _`Marcin Sztolcman`: https://github.com/msztolcman
diff --git a/pgcli/main.py b/pgcli/main.py
index 1fc17212..846b41cc 100644
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -118,12 +118,13 @@ class PGCli(object):
 
     def __init__(self, force_passwd_prompt=False, never_passwd_prompt=False,
                  pgexecute=None, pgclirc_file=None, row_limit=None,
-                 single_connection=False, less_chatty=None, prompt=None,
+                 single_connection=False, less_chatty=None, prompt=None, prompt_dsn=None,
                  auto_vertical_output=False):
 
         self.force_passwd_prompt = force_passwd_prompt
         self.never_passwd_prompt = never_passwd_prompt
         self.pgexecute = pgexecute
+        self.dsn_alias = None
 
         # Load config.
         c = self.config = get_config(pgclirc_file)
@@ -156,6 +157,7 @@ class PGCli(object):
         self.less_chatty = bool(less_chatty) or c['main'].as_bool('less_chatty')
         self.null_string = c['main'].get('null_string', '<null>')
         self.prompt_format = prompt if prompt is not None else c['main'].get('prompt', self.default_prompt)
+        self.prompt_dsn_format = prompt_dsn
         self.on_error = c['main']['on_error'].upper()
         self.decimal_format = c['data_formats']['decimal']
         self.float_format = c['data_formats']['float']
@@ -562,10 +564,17 @@ class PGCli(object):
             set_vi_mode_enabled=set_vi_mode)
 
         def prompt_tokens(_):
-            prompt = self.get_prompt(self.prompt_format)
-            if (self.prompt_format == self.default_prompt and
-               len(prompt) > self.max_len_prompt):
+            if self.dsn_alias and self.prompt_dsn_format is not None:
+                prompt_format = self.prompt_dsn_format
+            else:
+                prompt_format = self.prompt_format
+
+            prompt = self.get_prompt(prompt_format)
+
+            if (prompt_format == self.default_prompt and
+                    len(prompt) > self.max_len_prompt):
                 prompt = self.get_prompt('\\d> ')
+
             return [(Token.Prompt, prompt)]
 
         def get_continuation_tokens(cli, width):
@@ -776,6 +785,8 @@ class PGCli(object):
                 Document(text=text, cursor_position=cursor_positition), None)
 
     def get_prompt(self, string):
+        # should be before replacing \\d
+        string = string.replace('\\dsn_alias', self.dsn_alias or '')
         string = string.replace('\\t', self.now.strftime('%x %X'))
         string = string.replace('\\u', self.pgexecute.user or '(none)')
         string = string.replace('\\h', self.pgexecute.host or '(none)')
@@ -819,6 +830,7 @@ class PGCli(object):
         default=False,
         help='Skip intro on startup and goodbye on exit.')
 @click.option('--prompt', help='Prompt format (Default: "\\u@\\h:\\d> ").')
+@click.option('--prompt-dsn', help='Prompt format for connections using DSN aliases (Default: "\\u@\\h:\\d> ").')
 @click.option('-l', '--list', 'list_databases', is_flag=True, help='list '
               'available databases, then exit.')
 @click.option('--auto-vertical-output', is_flag=True,
@@ -827,7 +839,7 @@ class PGCli(object):
 @click.argument('username', default=lambda: None, envvar='PGUSER', nargs=1)
 def cli(database, username_opt, host, port, prompt_passwd, never_prompt,
         single_connection, dbname, username, version, pgclirc, dsn, row_limit,
-        less_chatty, prompt, list_databases, auto_vertical_output):
+        less_chatty, prompt, prompt_dsn, list_databases, auto_vertical_output):
 
     if version:
         print('Version:', __version__)
@@ -851,7 +863,7 @@ def cli(database, username_opt, host, port, prompt_passwd, never_prompt,
 
     pgcli = PGCli(prompt_passwd, never_prompt, pgclirc_file=pgclirc,
                   row_limit=row_limit, single_connection=single_connection,
-                  less_chatty=less_chatty, prompt=prompt,
+                  less_chatty=less_chatty, prompt=prompt, prompt_dsn=prompt_dsn,
                   auto_vertical_output=auto_vertical_output)
 
     # Choose which ever one has a valid value.
@@ -872,6 +884,7 @@ def cli(database, username_opt, host, port, prompt_passwd, never_prompt,
                  err=True, fg='red')
             exit(1)
         pgcli.connect_uri(dsn_config)
+        pgcli.dsn_alias = dsn
     elif '://' in database:
         pgcli.connect_uri(database)
     elif "=" in database:
diff --git a/pgcli/pgclirc b/pgcli/pgclirc
index 8469d918..3ed4662f 100644
--- a/pgcli/pgclirc
+++ b/pgcli/pgclirc
@@ -111,10 +111,15 @@ row_limit = 1000
 less_chatty = False
 
 # Postgres prompt
+# \t - Current date and time
 # \u - Username
 # \h - Hostname of the server
 # \d - Database name
+# \p - Database port
+# \i - Postgres PID
+# \# - "@" sign if logged in as superuser, '>' in other case
 # \n - Newline
+# \dsn_alias - name of dsn alias if -D option is used (empty otherwise)
 prompt = '\u@\h:\d> '
 
 # Number of lines to reserve for the suggestion menu