feat(commands): reset_db for Postgres backend

- connects to POSTGRES_HOST as an elevated POSTGRES_USER
- drops and recreates DJANGO_USER
- drops and recreates DJANGO_DB, giving ownership to DJANGO_USER
  with some other minor configuration

override database engine for tests where we continue to use SQLite

Author: lalver1

bin/reset_db.sh

@@ -5,12 +5,7 @@ set -ex
 DB_RESET="${DJANGO_DB_RESET:-true}"
 
 if [[ $DB_RESET = true ]]; then
-    # construct the path to the database file from environment or default
-    DB_DIR="${DJANGO_STORAGE_DIR:-.}"
-    DB_FILE="${DJANGO_DB_FILE:-django.db}"
-    DB_PATH="${DB_DIR}/${DB_FILE}"
-
-    rm -f "${DB_PATH}"
+    python manage.py reset_db
 
     # run database migrations and other initialization
     bin/init.sh

pems/core/management/__init__.py

@@ -0,0 +1,75 @@
+import os
+
+import psycopg
+
+from django.core.management.base import BaseCommand
+from django.db import connection
+
+
+class Command(BaseCommand):
+    help = "Completely resets the database (DESTRUCTIVE)."
+
+    def admin_connection(self) -> psycopg.Connection:
+        db_host = connection.settings_dict["HOST"]
+        db_port = connection.settings_dict["PORT"]
+
+        postgres_db = os.environ.get("POSTGRES_DB", "postgres")
+        admin_name = os.environ.get("POSTGRES_USER", "postgres")
+        admin_password = os.environ.get("POSTGRES_PASSWORD")
+
+        if not admin_password:
+            self.stderr.write(self.style.ERROR("POSTGRES_PASSWORD environment variable not set."))
+            return
+
+        return psycopg.connect(
+            host=db_host,
+            port=db_port,
+            user=admin_name,
+            password=admin_password,
+            dbname=postgres_db,
+            # Execute SQL commands immediately
+            autocommit=True,
+        )
+
+    def handle(self, *args, **options):
+        db_name = connection.settings_dict["NAME"]
+        db_user = connection.settings_dict["USER"]
+        db_password = connection.settings_dict["PASSWORD"]
+
+        if not db_password:
+            self.stderr.write(self.style.ERROR("DJANGO_DB_PASSWORD environment variable not set."))
+            return
+
+        try:
+            with self.admin_connection() as conn:
+                cursor = conn.cursor()
+                self.stdout.write(self.style.WARNING("Attempting database reset..."))
+
+                # Revoke existing connections
+                cursor.execute(
+                    """
+                    SELECT pg_terminate_backend(pg_stat_activity.pid)
+                    FROM pg_stat_activity
+                    WHERE datname = %s AND pg_stat_activity.pid <> pg_backend_pid()
+                    """,
+                    [db_name],
+                )
+                self.stdout.write(self.style.SUCCESS(f"Terminated existing connections to '{db_name}'."))
+
+                # Drop database
+                cursor.execute(f"DROP DATABASE IF EXISTS {db_name}")
+                self.stdout.write(self.style.SUCCESS(f"Database '{db_name}' dropped."))
+
+                # Create Django user
+                cursor.execute(f"DROP USER IF EXISTS {db_user}")
+                cursor.execute(f"CREATE USER {db_user} WITH PASSWORD '{db_password}'")
+                self.stdout.write(self.style.SUCCESS(f"Django user '{db_user}' created."))
+
+                # Create Django database with Django user as owner
+                cursor.execute(f"CREATE DATABASE {db_name} WITH OWNER {db_user} ENCODING 'UTF-8'")
+                self.stdout.write(self.style.SUCCESS(f"Database '{db_name}' created and owned by '{db_user}'."))
+
+                self.stdout.write(self.style.SUCCESS("Database reset and user setup complete."))
+
+        except Exception as e:
+            self.stderr.write(self.style.ERROR(f"Error during database reset: {e}"))

pems/core/management/commands/__init__.py

@@ -48,7 +48,7 @@ preserve_blank_lines = true
 use_gitignore = true
 
 [tool.pytest.ini_options]
-DJANGO_SETTINGS_MODULE = "pems.settings"
+DJANGO_SETTINGS_MODULE = "tests.pytest.settings"
 
 [tool.setuptools.packages.find]
 include = ["pems*"]

pems/core/management/commands/reset_db.py

@@ -0,0 +1,111 @@
+import pytest
+
+from pems.core.management.commands import reset_db
+
+
+@pytest.fixture
+def command(mocker):
+    cmd = reset_db.Command()
+    # mocking command I/O for tests
+    cmd.stderr = mocker.Mock()
+    cmd.stdout = mocker.Mock()
+    cmd.style.ERROR = str
+    cmd.style.SUCCESS = str
+    cmd.style.WARNING = str
+    return cmd
+
+
+@pytest.fixture(autouse=True)
+def db_settings(settings):
+    settings.DATABASES["default"]["NAME"] = "test_db"
+    settings.DATABASES["default"]["USER"] = "test_user"
+    settings.DATABASES["default"]["PASSWORD"] = "test_password"
+    settings.DATABASES["default"]["HOST"] = "test_host"
+    settings.DATABASES["default"]["PORT"] = "1234"
+    return settings
+
+
+@pytest.fixture
+def mock_connect(mocker):
+    mock = mocker.patch.object(reset_db.psycopg, "connect", mocker.Mock())
+    return mock
+
+
+@pytest.fixture
+def mock_cursor(mocker):
+    return mocker.Mock()
+
+
+@pytest.fixture
+def mock_conn(mocker, mock_cursor):
+    mock = mocker.MagicMock()
+    # fake context manager support for `with manager() as ctx:`
+    mock.__enter__.return_value = mock
+    mock.cursor.return_value = mock_cursor
+    return mock
+
+
+@pytest.fixture
+def mock_admin_connection(mocker, mock_conn):
+    return mocker.patch.object(reset_db.Command, "admin_connection", return_value=mock_conn)
+
+
+@pytest.mark.django_db
+class TestCommand:
+    @pytest.mark.parametrize(
+        "env_vars,should_connect",
+        [
+            ({"POSTGRES_DB": "test_db", "POSTGRES_USER": "test_user", "POSTGRES_PASSWORD": "test_pass"}, True),
+            ({"POSTGRES_DB": "test_db", "POSTGRES_USER": "test_user", "POSTGRES_PASSWORD": ""}, False),
+        ],
+    )
+    def test_admin_connection(self, command: reset_db.Command, monkeypatch, mock_connect, env_vars, should_connect):
+        for key, value in env_vars.items():
+            monkeypatch.setenv(key, value)
+
+        result = command.admin_connection()
+
+        if should_connect:
+            mock_connect.assert_called_once_with(
+                host="test_host", port="1234", user="test_user", password="test_pass", dbname="test_db", autocommit=True
+            )
+            assert result is not None
+        else:
+            assert result is None
+
+    def test_handle_success(self, command: reset_db.Command, mocker, mock_admin_connection, mock_cursor):
+        mocker.patch("os.environ.get", return_value="test_password")
+
+        command.handle()
+
+        mock_admin_connection.assert_called_once()
+        assert mock_cursor.execute.call_count == 5
+
+        # Verify success messages were logged
+        expected_messages = [
+            "Attempting database reset...",
+            "Terminated existing connections to 'test_db'.",
+            "Database 'test_db' dropped.",
+            "Django user 'test_user' created.",
+            "Database 'test_db' created and owned by 'test_user'.",
+            "Database reset and user setup complete.",
+        ]
+
+        for message in expected_messages:
+            command.stdout.write.assert_any_call(message)
+
+    def test_handle_no_db_password(self, command: reset_db.Command, db_settings, mock_admin_connection):
+        db_settings.DATABASES["default"]["PASSWORD"] = None
+
+        command.handle()
+
+        mock_admin_connection.assert_not_called()
+        command.stderr.write.assert_any_call("DJANGO_DB_PASSWORD environment variable not set.")
+
+    def test_handle_admin_connection_error(self, command: reset_db.Command, mock_admin_connection):
+        mock_admin_connection.side_effect = Exception("Admin connection failed.")
+
+        command.handle()
+
+        mock_admin_connection.assert_called_once()
+        command.stderr.write.assert_any_call("Error during database reset: Admin connection failed.")

pyproject.toml

@@ -0,0 +1,8 @@
+from pems.settings import *  # noqa: F401, F403
+
+DATABASES = {
+    "default": {
+        "ENGINE": "django.db.backends.sqlite3",
+        "NAME": "test",
+    }
+}