From ff2e810358fcf0bef7a52b6d67ee8bd2c018d595 Mon Sep 17 00:00:00 2001
From: Simon Percivall <percivall@gmail.com>
Date: Thu, 6 Aug 2015 16:38:45 +0200
Subject: [PATCH] Add `dbshell_plus` management command.

---
 lib/django_dbshell_plus/__init__.py           |  2 +-
 .../management/__init__.py                    |  0
 .../management/commands/__init__.py           |  0
 .../management/commands/dbshell_plus.py       | 67 +++++++++++++++++++
 4 files changed, 68 insertions(+), 1 deletion(-)
 create mode 100644 lib/django_dbshell_plus/management/__init__.py
 create mode 100644 lib/django_dbshell_plus/management/commands/__init__.py
 create mode 100644 lib/django_dbshell_plus/management/commands/dbshell_plus.py

diff --git a/lib/django_dbshell_plus/__init__.py b/lib/django_dbshell_plus/__init__.py
index 17fc974..c03f526 100644
--- a/lib/django_dbshell_plus/__init__.py
+++ b/lib/django_dbshell_plus/__init__.py
@@ -1,2 +1,2 @@
 # coding: utf-8
-__version__ = '1.0'
+__version__ = '1.0.1'
diff --git a/lib/django_dbshell_plus/management/__init__.py b/lib/django_dbshell_plus/management/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/lib/django_dbshell_plus/management/commands/__init__.py b/lib/django_dbshell_plus/management/commands/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/lib/django_dbshell_plus/management/commands/dbshell_plus.py b/lib/django_dbshell_plus/management/commands/dbshell_plus.py
new file mode 100644
index 0000000..720f800
--- /dev/null
+++ b/lib/django_dbshell_plus/management/commands/dbshell_plus.py
@@ -0,0 +1,67 @@
+import errno
+import subprocess
+
+from django.db import connections
+from django.core.management.commands import dbshell
+
+
+class Command(dbshell.Command):
+
+    def handle(self, **options):
+        connection = connections[options.get('database')]
+
+        dbclis = {
+            'postgresql': 'pgcli',
+            'mysql': 'mycli'
+        }
+
+        cmd = dbclis.get(connection.vendor)
+        if cmd:
+            try:
+                getattr(self, cmd)(connection)
+                return
+            except OSError, e:
+                if e.errno != errno.ENOENT:
+                    self.stderr.write("Could not start %s: %s" % (cmd, str(e)))
+
+        super(Command, self).handle(**options)
+
+    def pgcli(self, connection):
+        # argument code copied from Django
+        settings_dict = connection.settings_dict
+        args = ['pgcli']
+        if settings_dict['USER']:
+            args += ["-U", settings_dict['USER']]
+        if settings_dict['HOST']:
+            args.extend(["-h", settings_dict['HOST']])
+        if settings_dict['PORT']:
+            args.extend(["-p", str(settings_dict['PORT'])])
+        args += [settings_dict['NAME']]
+
+        subprocess.call(args)
+
+    def mycli(self, connection):
+        # argument code copied from Django
+        settings_dict = connection.settings_dict
+        args = ['mycli']
+        db = settings_dict['OPTIONS'].get('db', settings_dict['NAME'])
+        user = settings_dict['OPTIONS'].get('user', settings_dict['USER'])
+        passwd = settings_dict['OPTIONS'].get('passwd', settings_dict['PASSWORD'])
+        host = settings_dict['OPTIONS'].get('host', settings_dict['HOST'])
+        port = settings_dict['OPTIONS'].get('port', settings_dict['PORT'])
+
+        if user:
+            args += ["--user=%s" % user]
+        if passwd:
+            args += ["--pass=%s" % passwd]
+        if host:
+            if '/' in host:
+                args += ["--socket=%s" % host]
+            else:
+                args += ["--host=%s" % host]
+        if port:
+            args += ["--port=%s" % port]
+        if db:
+            args += [db]
+
+        subprocess.call(args)