From 557c9cddb24c6b8e3a37646498721c75ab5a4bd3 Mon Sep 17 00:00:00 2001
From: Jonathan Tsai <hello@jontsai.com>
Date: Tue, 28 May 2019 04:33:58 +0000
Subject: [PATCH] try mycli/pgcli in order of: virtualenv > system > regular
 mysql/psql

---
 AUTHORS.rst                                   |  1 +
 .../management/commands/dbshell_plus.py       | 41 ++++++++++++++++---
 2 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/AUTHORS.rst b/AUTHORS.rst
index 4a04b19..c6471d0 100644
--- a/AUTHORS.rst
+++ b/AUTHORS.rst
@@ -3,3 +3,4 @@ Credits
 =======
 
 * Simon Percivall <percivall@gmail.com>
+* Jonathan Tsai <hello@jontsai.com>
diff --git a/lib/django_dbshell_plus/management/commands/dbshell_plus.py b/lib/django_dbshell_plus/management/commands/dbshell_plus.py
index 720f800..be33068 100644
--- a/lib/django_dbshell_plus/management/commands/dbshell_plus.py
+++ b/lib/django_dbshell_plus/management/commands/dbshell_plus.py
@@ -1,5 +1,6 @@
 import errno
 import subprocess
+import sys
 
 from django.db import connections
 from django.core.management.commands import dbshell
@@ -18,18 +19,46 @@ def handle(self, **options):
         cmd = dbclis.get(connection.vendor)
         if cmd:
             try:
+                # attempt to use mycli/pgcli
                 getattr(self, cmd)(connection)
                 return
             except OSError, e:
-                if e.errno != errno.ENOENT:
-                    self.stderr.write("Could not start %s: %s" % (cmd, str(e)))
+                if self._is_virtualenv:
+                    try:
+                        # retry without explicitly without virtualenv
+                        getattr(self, cmd)(connection, ignore_virtualenv=True)
+                        return
+                    except OSError, e:
+                        if e.errno != errno.ENOENT:
+                            self.stderr.write("Could not start %s: %s" % (cmd, str(e)))
+                else:
+                    if e.errno != errno.ENOENT:
+                        self.stderr.write("Could not start %s: %s" % (cmd, str(e)))
 
+        # default to system mysql/psql
         super(Command, self).handle(**options)
 
-    def pgcli(self, connection):
+    @property
+    def _is_virtualenv(self):
+        # sys.real_prefix is only set if inside virtualenv
+        return hasattr(sys, 'real_prefix')
+
+    @property
+    def _python_path(self):
+        path = '{}/bin/'.format(sys.prefix) if self._is_virtualenv else ''
+        return path
+
+    def _get_cli_command(self, cli, ignore_virtualenv=False):
+        cli_command = '{}{}'.format(
+            '' if ignore_virtualenv else self._python_path,
+            cli, # 'pgcli' or 'mycli'
+        )
+        return cli_command
+
+    def pgcli(self, connection, ignore_virtualenv=False):
         # argument code copied from Django
         settings_dict = connection.settings_dict
-        args = ['pgcli']
+        args = [self._get_cli_command('pgcli', ignore_virtualenv=ignore_virtualenv)]
         if settings_dict['USER']:
             args += ["-U", settings_dict['USER']]
         if settings_dict['HOST']:
@@ -40,10 +69,10 @@ def pgcli(self, connection):
 
         subprocess.call(args)
 
-    def mycli(self, connection):
+    def mycli(self, connection, ignore_virtualenv=False):
         # argument code copied from Django
         settings_dict = connection.settings_dict
-        args = ['mycli']
+        args = [self._get_cli_command('mycli', ignore_virtualenv=ignore_virtualenv)]
         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'])