Repository: django-extensions/django-extensions

Issue:
`drop_test_database` not compatible with psycopg3
It appears that the `drop_test_database` command is hardcoded for psycopg2 and will break with psycopg3 (supported on Django 4.2+):

Python: 3.11
Django: 4.2
Postgres: 15.0
psycopg[binary]: 3.18

```bash
You have requested to drop all test databases.
This will IRREVERSIBLY DESTROY
ALL data in the database "test_postgres"
and all cloned test databases generated via
the "--parallel" flag (these are sequentially
named "test_postgres_1", "test_postgres_2", etc.).
Are you sure you want to do this?

Type 'yes' to continue, or 'no' to cancel: yes
Traceback (most recent call last):
  File "./manage.py", line 26, in <module>
    main()
  File "./manage.py", line 22, in main
    execute_from_command_line(sys.argv)
  File "my-path/lib/python3.11/site-packages/django/core/management/__init__.py", line 442, in execute_from_command_line
    utility.execute()
  File "/my-path/lib/python3.11/site-packages/django/core/management/__init__.py", line 436, in execute
    self.fetch_command(subcommand).run_from_argv(self.argv)
  File "/my-path/lib/python3.11/site-packages/django/core/management/base.py", line 412, in run_from_argv
    self.execute(*args, **cmd_options)
  File "/my-path/lib/python3.11/site-packages/django/core/management/base.py", line 458, in execute
    output = self.handle(*args, **options)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/my-path/lib/python3.11/site-packages/django_extensions/management/utils.py", line 62, in inner
    ret = func(self, *args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/my-path/lib/python3.11/site-packages/django_extensions/management/commands/drop_test_database.py", line 170, in handle
    import psycopg2 as Database  # NOQA

ModuleNotFoundError: No module named 'psycopg2'

```


Discussion:
Same with `reset_db.py`
Is there a workaround at all?

```text

Traceback (most recent call last):
  File "/app/manage.py", line 29, in <module>
    execute_from_command_line(sys.argv)
  File "/usr/local/lib/python3.11/site-packages/django/core/management/__init__.py", line 442, in execute_from_command_line
    utility.execute()
  File "/usr/local/lib/python3.11/site-packages/django/core/management/__init__.py", line 436, in execute
    self.fetch_command(subcommand).run_from_argv(self.argv)
  File "/usr/local/lib/python3.11/site-packages/django/core/management/base.py", line 412, in run_from_argv
    self.execute(*args, **cmd_options)
  File "/usr/local/lib/python3.11/site-packages/django/core/management/base.py", line 458, in execute
    output = self.handle(*args, **options)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/django_extensions/management/utils.py", line 62, in inner
    ret = func(self, *args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/django_extensions/management/commands/reset_db.py", line 144, in handle
    import psycopg2 as Database  # NOQA
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^
ModuleNotFoundError: No module named 'psycopg2'
```

Resolving patch:
diff --git a/django_extensions/management/commands/reset_db.py b/django_extensions/management/commands/reset_db.py
--- a/django_extensions/management/commands/reset_db.py
+++ b/django_extensions/management/commands/reset_db.py
@@ -4,6 +4,7 @@
 
 originally from http://www.djangosnippets.org/snippets/828/ by dnordberg
 """
+import importlib.util
 import os
 import logging
 import warnings
@@ -141,7 +142,11 @@ def handle(self, *args, **options):
             logging.info('Executing... "%s"', create_query)
             connection.query(create_query.strip())
         elif engine in POSTGRESQL_ENGINES:
-            import psycopg2 as Database  # NOQA
+            has_psycopg3 = importlib.util.find_spec("psycopg")
+            if has_psycopg3:
+                import psycopg as Database  # NOQA
+            else:
+                import psycopg2 as Database  # NOQA
 
             conn_params = {'database': 'template1'}
             if user:


Test patch:
diff --git a/django_extensions/management/commands/drop_test_database.py b/django_extensions/management/commands/drop_test_database.py
--- a/django_extensions/management/commands/drop_test_database.py
+++ b/django_extensions/management/commands/drop_test_database.py
@@ -1,4 +1,5 @@
 # -*- coding: utf-8 -*-
+import importlib.util
 from itertools import count
 import os
 import logging
@@ -167,7 +168,11 @@ def format_filename(name, number):
                 cursor.execute(drop_query)
 
         elif engine in POSTGRESQL_ENGINES:
-            import psycopg2 as Database  # NOQA
+            has_psycopg3 = importlib.util.find_spec("psycopg")
+            if has_psycopg3:
+                import psycopg as Database  # NOQA
+            else:
+                import psycopg2 as Database  # NOQA
 
             conn_params = {'database': 'template1'}
             if user:
diff --git a/tests/management/commands/test_drop_test_database.py b/tests/management/commands/test_drop_test_database.py
--- a/tests/management/commands/test_drop_test_database.py
+++ b/tests/management/commands/test_drop_test_database.py
@@ -1,4 +1,5 @@
 # -*- coding: utf-8 -*-
+import importlib.util
 from io import StringIO
 from unittest.mock import MagicMock, Mock, PropertyMock, call, patch
 
@@ -251,7 +252,11 @@ def test_postgresql_should_drop_database(self, m_stdout):
         # Indicate that no clone databases exist
         type(m_cursor).rowcount = PropertyMock(side_effect=(1, 0))
 
-        with patch.dict("sys.modules", psycopg2=m_database):
+        mock_kwargs = {"psycopg2": m_database}
+        has_psycopg3 = importlib.util.find_spec("psycopg") is not None
+        if has_psycopg3:
+            mock_kwargs = {"psycopg": m_database}
+        with patch.dict("sys.modules", **mock_kwargs):
             call_command('drop_test_database', '--noinput', verbosity=2)
 
         with self.subTest('Should check for and remove test database names until failure'):
@@ -277,7 +282,11 @@ def test_postgresql_should_drop_all_existing_cloned_databases(self):
         # Indicate that clone databases exist up to test_test_2
         type(m_cursor).rowcount = PropertyMock(side_effect=(1, 1, 1, 0))
 
-        with patch.dict("sys.modules", psycopg2=m_database):
+        mock_kwargs = {"psycopg2": m_database}
+        has_psycopg3 = importlib.util.find_spec("psycopg") is not None
+        if has_psycopg3:
+            mock_kwargs = {"psycopg": m_database}
+        with patch.dict("sys.modules", **mock_kwargs):
             call_command('drop_test_database', '--noinput')
 
         exists_query = "SELECT datname FROM pg_catalog.pg_database WHERE datname="
@@ -303,7 +312,11 @@ def test_postgresql_should_not_print_Reset_successful_when_exception_occured(sel
         m_cursor.execute.side_effect = m_database.ProgrammingError
         m_database.connect.return_value.cursor.return_value = m_cursor
 
-        with patch.dict("sys.modules", psycopg2=m_database):
+        mock_kwargs = {"psycopg2": m_database}
+        has_psycopg3 = importlib.util.find_spec("psycopg") is not None
+        if has_psycopg3:
+            mock_kwargs = {"psycopg": m_database}
+        with patch.dict("sys.modules", **mock_kwargs):
             call_command('drop_test_database', '--noinput', verbosity=2)
 
         self.assertNotIn("Reset successful.", m_stdout.getvalue())
diff --git a/tests/management/commands/test_pipchecker.py b/tests/management/commands/test_pipchecker.py
--- a/tests/management/commands/test_pipchecker.py
+++ b/tests/management/commands/test_pipchecker.py
@@ -140,6 +140,7 @@ def test_pipchecker_with_long_up_to_date_requirements(self):
         f.write('Pillow')
         f.write('pluggy')
         f.write('psycopg2-binary')
+        f.write('psycopg[binary,pool]')
         f.write('py')
         f.write('pyparsing')
         f.write('pytest')
diff --git a/tests/management/commands/test_reset_db.py b/tests/management/commands/test_reset_db.py
--- a/tests/management/commands/test_reset_db.py
+++ b/tests/management/commands/test_reset_db.py
@@ -1,4 +1,5 @@
 # -*- coding: utf-8 -*-
+import importlib.util
 import os
 from io import StringIO
 
@@ -166,7 +167,11 @@ def test_should_drop_and_create_database_and_print_success_messsage(self, m_stdo
             mock.call('CREATE DATABASE "test_db" WITH OWNER = "foo"  ENCODING = \'UTF8\';'),
         ]
 
-        with mock.patch.dict("sys.modules", psycopg2=m_database):
+        mock_kwargs = {"psycopg2": m_database}
+        has_psycopg3 = importlib.util.find_spec("psycopg") is not None
+        if has_psycopg3:
+            mock_kwargs = {"psycopg": m_database}
+        with mock.patch.dict("sys.modules", **mock_kwargs):
             call_command('reset_db', '--noinput', verbosity=2)
 
         m_database.connect.assert_called_once_with(database='template1', host='127.0.0.1', password='bar', port='5432', user='foo')
@@ -186,7 +191,11 @@ def test_should_drop_create_database_close_sessions_and_print_success_messsage(s
             mock.call('CREATE DATABASE "test_db" WITH OWNER = "foo"  ENCODING = \'UTF8\' TABLESPACE = TEST_TABLESPACE;'),
         ]
 
-        with mock.patch.dict("sys.modules", psycopg2=m_database):
+        mock_kwargs = {"psycopg2": m_database}
+        has_psycopg3 = importlib.util.find_spec("psycopg") is not None
+        if has_psycopg3:
+            mock_kwargs = {"psycopg": m_database}
+        with mock.patch.dict("sys.modules", **mock_kwargs):
             call_command('reset_db', '--noinput', '--close-sessions', verbosity=2)
 
         m_database.connect.assert_called_once_with(database='template1', host='127.0.0.1', password='bar', port='5432', user='foo')