Skip to content

Commit 557c9cd

Browse files
committed
try mycli/pgcli in order of: virtualenv > system > regular mysql/psql
1 parent af6a47d commit 557c9cd

File tree

2 files changed

+36
-6
lines changed

2 files changed

+36
-6
lines changed

AUTHORS.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,3 +3,4 @@ Credits
33
=======
44

55
* Simon Percivall <[email protected]>
6+
* Jonathan Tsai <[email protected]>

lib/django_dbshell_plus/management/commands/dbshell_plus.py

Lines changed: 35 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import errno
22
import subprocess
3+
import sys
34

45
from django.db import connections
56
from django.core.management.commands import dbshell
@@ -18,18 +19,46 @@ def handle(self, **options):
1819
cmd = dbclis.get(connection.vendor)
1920
if cmd:
2021
try:
22+
# attempt to use mycli/pgcli
2123
getattr(self, cmd)(connection)
2224
return
2325
except OSError, e:
24-
if e.errno != errno.ENOENT:
25-
self.stderr.write("Could not start %s: %s" % (cmd, str(e)))
26+
if self._is_virtualenv:
27+
try:
28+
# retry without explicitly without virtualenv
29+
getattr(self, cmd)(connection, ignore_virtualenv=True)
30+
return
31+
except OSError, e:
32+
if e.errno != errno.ENOENT:
33+
self.stderr.write("Could not start %s: %s" % (cmd, str(e)))
34+
else:
35+
if e.errno != errno.ENOENT:
36+
self.stderr.write("Could not start %s: %s" % (cmd, str(e)))
2637

38+
# default to system mysql/psql
2739
super(Command, self).handle(**options)
2840

29-
def pgcli(self, connection):
41+
@property
42+
def _is_virtualenv(self):
43+
# sys.real_prefix is only set if inside virtualenv
44+
return hasattr(sys, 'real_prefix')
45+
46+
@property
47+
def _python_path(self):
48+
path = '{}/bin/'.format(sys.prefix) if self._is_virtualenv else ''
49+
return path
50+
51+
def _get_cli_command(self, cli, ignore_virtualenv=False):
52+
cli_command = '{}{}'.format(
53+
'' if ignore_virtualenv else self._python_path,
54+
cli, # 'pgcli' or 'mycli'
55+
)
56+
return cli_command
57+
58+
def pgcli(self, connection, ignore_virtualenv=False):
3059
# argument code copied from Django
3160
settings_dict = connection.settings_dict
32-
args = ['pgcli']
61+
args = [self._get_cli_command('pgcli', ignore_virtualenv=ignore_virtualenv)]
3362
if settings_dict['USER']:
3463
args += ["-U", settings_dict['USER']]
3564
if settings_dict['HOST']:
@@ -40,10 +69,10 @@ def pgcli(self, connection):
4069

4170
subprocess.call(args)
4271

43-
def mycli(self, connection):
72+
def mycli(self, connection, ignore_virtualenv=False):
4473
# argument code copied from Django
4574
settings_dict = connection.settings_dict
46-
args = ['mycli']
75+
args = [self._get_cli_command('mycli', ignore_virtualenv=ignore_virtualenv)]
4776
db = settings_dict['OPTIONS'].get('db', settings_dict['NAME'])
4877
user = settings_dict['OPTIONS'].get('user', settings_dict['USER'])
4978
passwd = settings_dict['OPTIONS'].get('passwd', settings_dict['PASSWORD'])

0 commit comments

Comments
 (0)