From eab0011f3174404374bb166308728f17ebe05991 Mon Sep 17 00:00:00 2001
From: Amjith Ramanujam <amjith@newrelic.com>
Date: Sun, 7 Dec 2014 22:30:24 -0800
Subject: [PATCH] Implement change database for \c and use.

---
 TODO               |  3 +++
 pgcli/main.py      |  7 ++++---
 pgcli/pgexecute.py | 27 +++++++++++++++++++++++++++
 pgcli/pgspecial.py |  4 ----
 4 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/TODO b/TODO
index a8135054..74ec3a0d 100644
--- a/TODO
+++ b/TODO
@@ -1,4 +1,7 @@
 * [] Vendor in tabulate.
+* [] Create a separate pgspecial package.
+* [] Vendor in pgspecial package.
+* [] Add \c command.
 * [] Add MySQL commands (use, describe etc)
 * [X] Add exit keyword.
 * [] Show only table sensitive columns.
diff --git a/pgcli/main.py b/pgcli/main.py
index ef047b68..cec004a2 100755
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -46,7 +46,7 @@ def cli(database, user, password, host, port):
     except Exception as e:
         click.secho(e.message, err=True, fg='red')
         exit(1)
-    layout = Layout(before_input=DefaultPrompt('%s> ' % database),
+    layout = Layout(before_input=DefaultPrompt('%s> ' % pgexecute.dbname),
             menus=[CompletionsMenu()],
             lexer=SqlLexer)
     completer = PGCompleter(config.getboolean('main', 'smart_completion'))
@@ -59,13 +59,14 @@ def cli(database, user, password, host, port):
 
     try:
         while True:
+            cli.layout.before_input = DefaultPrompt('%s> ' % pgexecute.dbname)
             document = cli.read_input(on_exit=AbortAction.RAISE_EXCEPTION)
 
             # The reason we check here instead of inside the pgexecute is
             # because we want to raise the Exit exception which will be caught
             # by the try/except block that wraps the pgexecute.run() statement.
-            if (document.text.strip() == 'exit'
-                    or document.text.strip() == 'quit'):
+            if (document.text.strip().lower() == 'exit'
+                    or document.text.strip().lower() == 'quit'):
                 raise Exit
             try:
                 rows, headers, status = pgexecute.run(document.text)
diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py
index 03b2d34f..b0e8654d 100644
--- a/pgcli/pgexecute.py
+++ b/pgcli/pgexecute.py
@@ -13,10 +13,37 @@ class PGExecute(object):
     def __init__(self, database, user, password, host, port):
         self.conn = psycopg2.connect(database=database, user=user,
                 password=password, host=host, port=port)
+        self.dbname = database
+        self.user = user
+        self.password = password
+        self.host = host
+        self.port = port
         self.conn.autocommit = True
 
+    @staticmethod
+    def parse_pattern(sql):
+        command, _, arg = sql.partition(' ')
+        verbose = '+' in command
+        return (command.strip(), verbose, arg.strip())
+
     def run(self, sql):
+
+        # Remove spaces, eol and semi-colons.
         sql = sql.strip()
+        sql = sql.rstrip(';')
+
+        # Check if the command is a \c or 'use'. This is a special exception
+        # that cannot be offloaded to `pgspecial` lib. Because we have to
+        # change the database connection that we're connected to.
+        if sql.startswith('\c') or sql.lower().startswith('use'):
+            dbname = self.parse_pattern(sql)[2]
+            self.conn = psycopg2.connect(database=dbname,
+                    user=self.user, password=self.password, host=self.host,
+                    port=self.port)
+            self.dbname = dbname
+            return (None, None, 'You are now connected to database "%s" as '
+                    'user "%s"' % (self.dbname, self.user))
+
         with self.conn.cursor() as cur:
             if sql in self.special_commands:
                 cur.execute(self.special_commands[sql])
diff --git a/pgcli/pgspecial.py b/pgcli/pgspecial.py
index a7953f95..9c70e328 100644
--- a/pgcli/pgspecial.py
+++ b/pgcli/pgspecial.py
@@ -675,14 +675,10 @@ def sql_name_pattern(pattern):
 
     return schema, relname
 
-
 if __name__ == '__main__':
     import psycopg2
     con = psycopg2.connect(database='misago_testforum')
     cur = con.cursor()
-    #print describe_table_details(cur, 'django_migrations', False)
-    #rows, headers, status = describe_table_details(cur, 'django_migrations', False)
-    #rows, headers, status = describe_table_details(cur, None, False)
     table = sys.argv[1]
     for rows, headers, status in describe_table_details(cur, table, False):
         print tabulate(rows, headers, tablefmt='psql')