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

13 files changed, 117 insertions, 21 deletions

diff --git a/django/core/management/commands/dbshell.py b/django/core/management/commands/dbshell.py
index 6782ef48f3..fb3caed12a 100644
--- a/django/core/management/commands/dbshell.py
+++ b/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 @@ class Command(BaseCommand):
             '--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 @@ class Command(BaseCommand):
                 '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,
+            )
diff --git a/django/db/backends/base/client.py b/django/db/backends/base/client.py
index f6a0cda35d..32901764aa 100644
--- a/django/db/backends/base/client.py
+++ b/django/db/backends/base/client.py
@@ -8,5 +8,5 @@ class BaseDatabaseClient:
         # 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')
diff --git a/django/db/backends/mysql/client.py b/django/db/backends/mysql/client.py
index a596a650b3..e5b6a59fb2 100644
--- a/django/db/backends/mysql/client.py
+++ b/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 @@ class DatabaseClient(BaseDatabaseClient):
             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)
diff --git a/django/db/backends/mysql/creation.py b/django/db/backends/mysql/creation.py
index 79d2959d5f..99372fd7ee 100644
--- a/django/db/backends/mysql/creation.py
+++ b/django/db/backends/mysql/creation.py
@@ -55,9 +55,9 @@ class DatabaseCreation(BaseDatabaseCreation):
         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:
diff --git a/django/db/backends/oracle/client.py b/django/db/backends/oracle/client.py
index 243c018d03..434481ba81 100644
--- a/django/db/backends/oracle/client.py
+++ b/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)
diff --git a/django/db/backends/postgresql/client.py b/django/db/backends/postgresql/client.py
index 0efe0d47f0..9d390b3807 100644
--- a/django/db/backends/postgresql/client.py
+++ b/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 @@ class DatabaseClient(BaseDatabaseClient):
         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 @@ class DatabaseClient(BaseDatabaseClient):
             # 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)
diff --git a/django/db/backends/sqlite3/client.py b/django/db/backends/sqlite3/client.py
index f6b7077a0c..a71005fd5b 100644
--- a/django/db/backends/sqlite3/client.py
+++ b/django/db/backends/sqlite3/client.py
@@ -6,9 +6,10 @@ from django.db.backends.base.client import BaseDatabaseClient
 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)
diff --git a/docs/ref/django-admin.txt b/docs/ref/django-admin.txt
index 4011da32fc..eb151193ae 100644
--- a/docs/ref/django-admin.txt
+++ b/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
diff --git a/docs/releases/3.1.txt b/docs/releases/3.1.txt
index dbf599fb48..384fa43dfb 100644
--- a/docs/releases/3.1.txt
+++ b/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
 --------------------------------
 
diff --git a/tests/dbshell/test_mysql.py b/tests/dbshell/test_mysql.py
index 6b60e2722e..91c1a6cd50 100644
--- a/tests/dbshell/test_mysql.py
+++ b/tests/dbshell/test_mysql.py
@@ -76,5 +76,23 @@ class MySqlDbshellCommandTestCase(SimpleTestCase):
                 },
             }))
 
-    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)
diff --git a/tests/dbshell/test_oracle.py b/tests/dbshell/test_oracle.py
index 431134e3ac..41fbc07455 100644
--- a/tests/dbshell/test_oracle.py
+++ b/tests/dbshell/test_oracle.py
@@ -8,17 +8,19 @@ from django.test import SimpleTestCase
 
 @skipUnless(connection.vendor == 'oracle', 'Oracle tests')
 class OracleDbshellTests(SimpleTestCase):
-    def _run_dbshell(self, rlwrap=False):
+    def _run_dbshell(self, rlwrap=False, parameters=None):
         """Run runshell command and capture its arguments."""
         def _mock_subprocess_run(*args, **kwargs):
             self.subprocess_args = list(*args)
             return CompletedProcess(self.subprocess_args, 0)
 
+        if parameters is None:
+            parameters = []
         client = DatabaseClient(connection)
         self.subprocess_args = None
         with mock.patch('subprocess.run', new=_mock_subprocess_run):
             with mock.patch('shutil.which', return_value='/usr/bin/rlwrap' if rlwrap else None):
-                client.runshell()
+                client.runshell(parameters)
         return self.subprocess_args
 
     def test_without_rlwrap(self):
@@ -32,3 +34,9 @@ class OracleDbshellTests(SimpleTestCase):
             self._run_dbshell(rlwrap=True),
             ['/usr/bin/rlwrap', 'sqlplus', '-L', connection._connect_string()],
         )
+
+    def test_parameters(self):
+        self.assertEqual(
+            self._run_dbshell(parameters=['-HELP']),
+            ['sqlplus', '-L', connection._connect_string(), '-HELP'],
+        )
diff --git a/tests/dbshell/test_postgresql.py b/tests/dbshell/test_postgresql.py
index 7047283bb3..6de60eaef2 100644
--- a/tests/dbshell/test_postgresql.py
+++ b/tests/dbshell/test_postgresql.py
@@ -9,7 +9,7 @@ from django.test import SimpleTestCase
 
 class PostgreSqlDbshellCommandTestCase(SimpleTestCase):
 
-    def _run_it(self, dbinfo):
+    def _run_it(self, dbinfo, parameters=None):
         """
         That function invokes the runshell command, while mocking
         subprocess.run(). It returns a 2-tuple with:
@@ -21,8 +21,11 @@ class PostgreSqlDbshellCommandTestCase(SimpleTestCase):
             # PostgreSQL environment variables.
             self.pg_env = {key: env[key] for key in env if key.startswith('PG')}
             return subprocess.CompletedProcess(self.subprocess_args, 0)
+
+        if parameters is None:
+            parameters = []
         with mock.patch('subprocess.run', new=_mock_subprocess_run):
-            DatabaseClient.runshell_db(dbinfo)
+            DatabaseClient.runshell_db(dbinfo, parameters)
         return self.subprocess_args, self.pg_env
 
     def test_basic(self):
@@ -104,6 +107,12 @@ class PostgreSqlDbshellCommandTestCase(SimpleTestCase):
             )
         )
 
+    def test_parameters(self):
+        self.assertEqual(
+            self._run_it({'database': 'dbname'}, ['--help']),
+            (['psql', 'dbname', '--help'], {}),
+        )
+
     def test_sigint_handler(self):
         """SIGINT is ignored in Python and passed to psql to abort queries."""
         def _mock_subprocess_run(*args, **kwargs):
@@ -114,6 +123,6 @@ class PostgreSqlDbshellCommandTestCase(SimpleTestCase):
         # The default handler isn't SIG_IGN.
         self.assertNotEqual(sigint_handler, signal.SIG_IGN)
         with mock.patch('subprocess.run', new=_mock_subprocess_run):
-            DatabaseClient.runshell_db({})
+            DatabaseClient.runshell_db({}, [])
         # dbshell restores the original handler.
         self.assertEqual(sigint_handler, signal.getsignal(signal.SIGINT))
diff --git a/tests/dbshell/test_sqlite.py b/tests/dbshell/test_sqlite.py
index 46d09895c8..c3b2b1e28d 100644
--- a/tests/dbshell/test_sqlite.py
+++ b/tests/dbshell/test_sqlite.py
@@ -9,15 +9,17 @@ from django.test import SimpleTestCase
 
 @skipUnless(connection.vendor == 'sqlite', 'SQLite tests.')
 class SqliteDbshellCommandTestCase(SimpleTestCase):
-    def _run_dbshell(self):
+    def _run_dbshell(self, parameters=None):
         """Run runshell command and capture its arguments."""
         def _mock_subprocess_run(*args, **kwargs):
             self.subprocess_args = list(*args)
             return CompletedProcess(self.subprocess_args, 0)
 
+        if parameters is None:
+            parameters = []
         client = DatabaseClient(connection)
         with mock.patch('subprocess.run', new=_mock_subprocess_run):
-            client.runshell()
+            client.runshell(parameters)
         return self.subprocess_args
 
     def test_path_name(self):
@@ -29,3 +31,13 @@ class SqliteDbshellCommandTestCase(SimpleTestCase):
                 self._run_dbshell(),
                 ['sqlite3', 'test.db.sqlite3'],
             )
+
+    def test_parameters(self):
+        with mock.patch.dict(
+            connection.settings_dict,
+            {'NAME': Path('test.db.sqlite3')},
+        ):
+            self.assertEqual(
+                self._run_dbshell(['-help']),
+                ['sqlite3', 'test.db.sqlite3', '-help'],
+            )