Merge pull request piccolo-orm#17 from piccolo-orm/sql_shell

added sql_shell command

Author: dantownsend

docs/src/piccolo/projects_and_apps/included_apps.rst

@@ -70,6 +70,20 @@ Launches an iPython shell, and automatically imports all of your registered
 
     piccolo shell run
 
+sql_shell
+~~~~~~~~~
+
+Launches a SQL shell (``psql`` or ``sqlite3`` depending on the engine), using
+the connection settings defined in ``piccolo_conf.py``. It's convenient if you
+need to run raw SQL queries on your database.
+
+.. code-block:: bash
+
+    piccolo sql_shell run
+
+For it to work, the underlying command needs to be on the path (i.e. ``psql``
+or ``sqlite3`` depending on which you're using).
+
 -------------------------------------------------------------------------------
 
 Optional includes

piccolo/apps/sql_shell/__init__.py

@@ -0,0 +1,65 @@
+import os
+import signal
+import subprocess
+import typing as t
+
+from piccolo.engine.postgres import PostgresEngine
+from piccolo.engine.sqlite import SQLiteEngine
+from piccolo.engine.finder import engine_finder
+
+if t.TYPE_CHECKING:
+    from piccolo.engine.base import Engine
+
+
+def run():
+    """
+    Launch the SQL shell for the configured engine. For Postgres
+    this will be psql, and for SQLite it will be sqlite3.
+    """
+    engine: t.Optional[Engine] = engine_finder()
+
+    if engine is None:
+        raise ValueError(
+            "Unable to find the engine - make sure piccolo_conf is on the "
+            "path."
+        )
+
+    # Heavily inspired by Django's dbshell command
+    if isinstance(engine, PostgresEngine):
+        engine: PostgresEngine = engine
+
+        args = ["psql"]
+
+        host = engine.config.get("host")
+        port = engine.config.get("port")
+        user = engine.config.get("user")
+        password = engine.config.get("password")
+        database = engine.config.get("database")
+
+        if user:
+            args += ["-U", user]
+        if host:
+            args += ["-h", host]
+        if port:
+            args += ["-p", str(port)]
+        args += [database]
+
+        sigint_handler = signal.getsignal(signal.SIGINT)
+        subprocess_env = os.environ.copy()
+        if password:
+            subprocess_env["PGPASSWORD"] = str(password)
+        try:
+            # Allow SIGINT to pass to psql to abort queries.
+            signal.signal(signal.SIGINT, signal.SIG_IGN)
+            print("Enter \\q to exit")
+            subprocess.run(args, check=True, env=subprocess_env)
+        finally:
+            # Restore the original SIGINT handler.
+            signal.signal(signal.SIGINT, sigint_handler)
+
+    elif isinstance(engine, SQLiteEngine):
+        engine: SQLiteEngine = engine
+        print("Enter .quit to exit")
+        subprocess.run(
+            ["sqlite3", engine.connection_kwargs.get("database")], check=True
+        )

piccolo/apps/sql_shell/commands/__init__.py

@@ -0,0 +1,7 @@
+from piccolo.conf.apps import AppConfig
+from .commands.run import run
+
+
+APP_CONFIG = AppConfig(
+    app_name="sql_shell", migrations_folder_path="", commands=[run]
+)

piccolo/apps/sql_shell/commands/run.py

@@ -11,6 +11,7 @@
 from piccolo.apps.playground.piccolo_app import APP_CONFIG as playground_config
 from piccolo.apps.project.piccolo_app import APP_CONFIG as project_config
 from piccolo.apps.shell.piccolo_app import APP_CONFIG as shell_config
+from piccolo.apps.sql_shell.piccolo_app import APP_CONFIG as sql_shell_config
 from piccolo.apps.user.piccolo_app import APP_CONFIG as user_config
 
 
@@ -53,6 +54,7 @@ def main():
         playground_config,
         project_config,
         shell_config,
+        sql_shell_config,
         user_config,
     ]:
         for command in _app_config.commands:

piccolo/apps/sql_shell/piccolo_app.py

@@ -0,0 +1,14 @@
+from unittest import TestCase
+from unittest.mock import patch, MagicMock
+
+from piccolo.apps.sql_shell.commands.run import run
+
+
+class TestRun(TestCase):
+    @patch("piccolo.apps.sql_shell.commands.run.subprocess")
+    def test_run(self, subprocess: MagicMock):
+        """
+        A simple test to make sure it executes without raising any exceptions.
+        """
+        run()
+        self.assertTrue(subprocess.run.called)