Fixed #29501 -- Allowed dbshell to pass options to underlying tool.

Author: adamchainz

django/core/management/commands/dbshell.py

@@ -1,3 +1,5 @@
+import subprocess
+
 from django.core.management.base import BaseCommand, CommandError
 from django.db import DEFAULT_DB_ALIAS, connections
 
@@ -15,11 +17,13 @@ def add_arguments(self, parser):
  '--database', default=DEFAULT_DB_ALIAS,
  help='Nominates a database onto which to open a shell. Defaults to the "default" database.',
  )
+ parameters = parser.add_argument_group('parameters', prefix_chars='--')
+ parameters.add_argument('parameters', nargs='*')
 
  def handle(self, **options):
  connection = connections[options['database']]
  try:
- connection.client.runshell()
+ connection.client.runshell(options['parameters'])
  except FileNotFoundError:
  # Note that we're assuming the FileNotFoundError relates to the
  # command missing. It could be raised for some other reason, in
@@ -29,3 +33,11 @@ def handle(self, **options):
  'You appear not to have the %r program installed or on your path.' %
  connection.client.executable_name
  )
+ except subprocess.CalledProcessError as e:
+ raise CommandError(
+ '"%s" returned non-zero exit status %s.' % (
+ ' '.join(e.cmd),
+ e.returncode,
+ ),
+ returncode=e.returncode,
+ )

django/db/backends/base/client.py

@@ -8,5 +8,5 @@ def __init__(self, connection):
  # connection is an instance of BaseDatabaseWrapper.
  self.connection = connection
 
- def runshell(self):
+ def runshell(self, parameters):
  raise NotImplementedError('subclasses of BaseDatabaseClient must provide a runshell() method')

django/db/backends/mysql/client.py

@@ -7,7 +7,7 @@ class DatabaseClient(BaseDatabaseClient):
  executable_name = 'mysql'
 
  @classmethod
- def settings_to_cmd_args(cls, settings_dict):
+ def settings_to_cmd_args(cls, settings_dict, parameters):
  args = [cls.executable_name]
  db = settings_dict['OPTIONS'].get('db', settings_dict['NAME'])
  user = settings_dict['OPTIONS'].get('user', settings_dict['USER'])
@@ -41,8 +41,9 @@ def settings_to_cmd_args(cls, settings_dict):
  args += ["--ssl-key=%s" % client_key]
  if db:
  args += [db]
+ args.extend(parameters)
  return args
 
- def runshell(self):
- args = DatabaseClient.settings_to_cmd_args(self.connection.settings_dict)
+ def runshell(self, parameters):
+ args = DatabaseClient.settings_to_cmd_args(self.connection.settings_dict, parameters)
  subprocess.run(args, check=True)

django/db/backends/mysql/creation.py

@@ -55,9 +55,9 @@ def _clone_test_db(self, suffix, verbosity, keepdb=False):
  self._clone_db(source_database_name, target_database_name)
 
  def _clone_db(self, source_database_name, target_database_name):
- dump_args = DatabaseClient.settings_to_cmd_args(self.connection.settings_dict)[1:]
+ dump_args = DatabaseClient.settings_to_cmd_args(self.connection.settings_dict, [])[1:]
  dump_cmd = ['mysqldump', *dump_args[:-1], '--routines', '--events', source_database_name]
- load_cmd = DatabaseClient.settings_to_cmd_args(self.connection.settings_dict)
+ load_cmd = DatabaseClient.settings_to_cmd_args(self.connection.settings_dict, [])
  load_cmd[-1] = target_database_name
 
  with subprocess.Popen(dump_cmd, stdout=subprocess.PIPE) as dump_proc:

django/db/backends/oracle/client.py

@@ -8,10 +8,11 @@ class DatabaseClient(BaseDatabaseClient):
  executable_name = 'sqlplus'
  wrapper_name = 'rlwrap'
 
- def runshell(self):
+ def runshell(self, parameters):
  conn_string = self.connection._connect_string()
  args = [self.executable_name, "-L", conn_string]
  wrapper_path = shutil.which(self.wrapper_name)
  if wrapper_path:
  args = [wrapper_path, *args]
+ args.extend(parameters)
  subprocess.run(args, check=True)

django/db/backends/postgresql/client.py

@@ -9,7 +9,7 @@ class DatabaseClient(BaseDatabaseClient):
  executable_name = 'psql'
 
  @classmethod
- def runshell_db(cls, conn_params):
+ def runshell_db(cls, conn_params, parameters):
  args = [cls.executable_name]
 
  host = conn_params.get('host', '')
@@ -29,6 +29,7 @@ def runshell_db(cls, conn_params):
  if port:
  args += ['-p', str(port)]
  args += [dbname]
+ args.extend(parameters)
 
  sigint_handler = signal.getsignal(signal.SIGINT)
  subprocess_env = os.environ.copy()
@@ -50,5 +51,5 @@ def runshell_db(cls, conn_params):
  # Restore the original SIGINT handler.
  signal.signal(signal.SIGINT, sigint_handler)
 
- def runshell(self):
- DatabaseClient.runshell_db(self.connection.get_connection_params())
+ def runshell(self, parameters):
+ self.runshell_db(self.connection.get_connection_params(), parameters)

django/db/backends/sqlite3/client.py

@@ -6,9 +6,10 @@
 class DatabaseClient(BaseDatabaseClient):
  executable_name = 'sqlite3'
 
- def runshell(self):
+ def runshell(self, parameters):
  # TODO: Remove str() when dropping support for PY37.
  # args parameter accepts path-like objects on Windows since Python 3.8.
  args = [self.executable_name,
  str(self.connection.settings_dict['NAME'])]
+ args.extend(parameters)
  subprocess.run(args, check=True)

docs/ref/django-admin.txt

@@ -226,6 +226,33 @@ program manually.
 
 Specifies the database onto which to open a shell. Defaults to ``default``.
 
+.. django-admin-option:: -- ARGUMENTS
+
+.. versionadded:: 3.1
+
+Any arguments following a ``--`` divider will be passed on to the underlying
+command-line client. For example, with PostgreSQL you can use the ``psql``
+command's ``-c`` flag to execute a raw SQL query directly:
+
+.. console::
+
+ $ django-admin dbshell -- -c 'select current_user'
+ current_user
+ --------------
+ postgres
+ (1 row)
+
+On MySQL/MariaDB, you can do this with the ``mysql`` command's ``-e`` flag:
+
+.. console::
+
+ $ django-admin dbshell -- -e "select user()"
+ +----------------------+
+ | user() |
+ +----------------------+
+ | djangonaut@localhost |
+ +----------------------+
+
 .. note::
 
  Be aware that not all options set in the :setting:`OPTIONS` part of your

docs/releases/3.1.txt

@@ -317,6 +317,9 @@ Management Commands
  :attr:`~django.core.management.CommandError` allows customizing the exit
  status for management commands.
 
+* The new :option:`dbshell -- ARGUMENTS <dbshell -->` option allows passing
+ extra arguments to the command-line client for the database.
+
 Migrations
 ~~~~~~~~~~
 
@@ -505,6 +508,9 @@ backends.
  yields a cursor and automatically closes the cursor and connection upon
  exiting the ``with`` statement.
 
+* ``DatabaseClient.runshell()`` now requires an additional ``parameters``
+ argument as a list of extra arguments to pass on to the command-line client.
+
 Dropped support for MariaDB 10.1
 --------------------------------
 

tests/dbshell/test_mysql.py

@@ -76,5 +76,23 @@ def test_ssl_certificate_is_added(self):
  },
  }))
 
- def get_command_line_arguments(self, connection_settings):
- return DatabaseClient.settings_to_cmd_args(connection_settings)
+ def test_parameters(self):
+ self.assertEqual(
+ ['mysql', 'somedbname', '--help'],
+ self.get_command_line_arguments(
+ {
+ 'NAME': 'somedbname',
+ 'USER': None,
+ 'PASSWORD': None,
+ 'HOST': None,
+ 'PORT': None,
+ 'OPTIONS': {},
+ },
+ ['--help'],
+ ),
+ )
+
+ def get_command_line_arguments(self, connection_settings, parameters=None):
+ if parameters is None:
+ parameters = []
+ return DatabaseClient.settings_to_cmd_args(connection_settings, parameters)