8000 Merge pull request #2 from django-extensions/main · ryanchen99/django-extensions@94261d7 · GitHub
[go: up one dir, main page]

Skip to content

Commit 94261d7

Browse files
authored
Merge pull request #2 from django-extensions/main
Add support for psycopg3 (django-extensions#1814)
2 parents f30471d + bc43793 commit 94261d7

File tree

7 files changed

+46
-8
lines changed

7 files changed

+46
-8
lines changed

django_extensions/management/commands/drop_test_database.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
# -*- coding: utf-8 -*-
2+
import importlib.util
23
from itertools import count
34
import os
45
import logging
@@ -167,7 +168,11 @@ def format_filename(name, number):
167168
cursor.execute(drop_query)
168169

169170
elif engine in POSTGRESQL_ENGINES:
170-
import psycopg2 as Database # NOQA
171+
has_psycopg3 = importlib.util.find_spec("psycopg")
172+
if has_psycopg3:
173+
import psycopg as Database # NOQA
174+
else:
175+
import psycopg2 as Database # NOQA
171176

172177
conn_params = {'database': 'template1'}
173178
if user:

django_extensions/management/commands/reset_db.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
55
originally from http://www.djangosnippets.org/snippets/828/ by dnordberg
66
"""
7+
import importlib.util
78
import os
89
import logging
910
import warnings
@@ -141,7 +142,11 @@ def handle(self, *args, **options):
141142
logging.info('Executing... "%s"', create_query)
142143
connection.query(create_query.strip())
143144
elif engine in POSTGRESQL_ENGINES:
144-
import psycopg2 as Database # NOQA
145+
has_psycopg3 = importlib.util.find_spec("psycopg")
146+
if has_psycopg3:
147+
import psycopg as Database # NOQA
148+
else:
149+
import psycopg2 as Database # NOQA
145150

146151
conn_params = {'database': 'template1'}
147152
if user:

docs/sqldsn.rst

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ Supported Databases
1010

1111
Currently the following databases are supported:
1212

13-
* PostgreSQL (psycopg2 or postgis)
13+
* PostgreSQL (psycopg2, psycopg3, or postgis)
1414
* Sqlite3
1515
* MySQL
1616

tests/management/commands/test_drop_test_database.py

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
# -*- coding: utf-8 -*-
2+
import importlib.util
23
from io import StringIO
34
from unittest.mock import MagicMock, Mock, PropertyMock, call, patch
45

@@ -251,7 +252,11 @@ def test_postgresql_should_drop_database(self, m_stdout):
251252
# Indicate that no clone databases exist
252253
type(m_cursor).rowcount = PropertyMock(side_effect=(1, 0))
253254

254-
with patch.dict("sys.modules", psycopg2=m_database):
255+
mock_kwargs = {"psycopg2": m_database}
256+
has_psycopg3 = importlib.util.find_spec("psycopg") is not None
257+
if has_psycopg3:
258+
mock_kwargs = {"psycopg": m_database}
259+
with patch.dict("sys.modules", **mock_kwargs):
255260
call_command('drop_test_database', '--noinput', verbosity=2)
256261

257262
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):
277282
# Indicate that clone databases exist up to test_test_2
278283
type(m_cursor).rowcount = PropertyMock(side_effect=(1, 1, 1, 0))
279284

280-
with patch.dict("sys.modules", psycopg2=m_database):
285+
mock_kwargs = {"psycopg2": m_database}
286+
has_psycopg3 = importlib.util.find_spec("psycopg") is not None
287+
if has_psycopg3:
288+
mock_kwargs = {"psycopg": m_database}
289+
with patch.dict("sys.modules", **mock_kwargs):
281290
call_command('drop_test_database', '--noinput')
282291

283292
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
303312
m_cursor.execute.side_effect = m_database.ProgrammingError
304313
m_database.connect.return_value.cursor.return_value = m_cursor
305314

306-
with patch.dict("sys.modules", psycopg2=m_database):
315+
mock_kwargs = {"psycopg2": m_database}
316+
has_psycopg3 = importlib.util.find_spec("psycopg") is not None
317+
if has_psycopg3:
318+
mock_kwargs = {"psycopg": m_database}
319+
with patch.dict("sys.modules", **mock_kwargs):
307320
call_command('drop_test_database', '--noinput', verbosity=2)
308321

309322
self.assertNotIn("Reset successful.", m_stdout.getvalue())

tests/management/commands/test_pipchecker.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,7 @@ def test_pipchecker_with_long_up_to_date_requirements(self):
140140
f.write('Pillow')
141141
f.write('pluggy')
142142
f.write('psycopg2-binary')
143+
f.write('psycopg[binary,pool]')
143144
f.write('py')
144145
f.write('pyparsing')
145146
f.write('pytest')

tests/management/commands/test_reset_db.py

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
# -*- coding: utf-8 -*-
2+
import importlib.util
23
import os
34
from io import StringIO
45

@@ -166,7 +167,11 @@ def test_should_drop_and_create_database_and_print_success_messsage(self, m_stdo
166167
mock.call('CREATE DATABASE "test_db" WITH OWNER = "foo" ENCODING = \'UTF8\';'),
167168
]
168169

169-
with mock.patch.dict("sys.modules", psycopg2=m_database):
170+
mock_kwargs = {"psycopg2": m_database}
171+
has_psycopg3 = importlib.util.find_spec("psycopg") is not None
172+
if has_psycopg3:
173+
mock_kwargs = {"psycopg": m_database}
174+
with mock.patch.dict("sys.modules", **mock_kwargs):
170175
call_command('reset_db', '--noinput', verbosity=2)
171176

172177
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
186191
mock.call('CREATE DATABASE "test_db" WITH OWNER = "foo" ENCODING = \'UTF8\' TABLESPACE = TEST_TABLESPACE;'),
187192
]
188193

189-
with mock.patch.dict("sys.modules", psycopg2=m_database):
194+
mock_kwargs = {"psycopg2": m_database}
195+
has_psycopg3 = importlib.util.find_spec("psycopg") is not None
196+
if has_psycopg3:
197+
mock_kwargs = {"psycopg": m_database}
198+
with mock.patch.dict("sys.modules", **mock_kwargs):
190199
call_command('reset_db', '--noinput', '--close-sessions', verbosity=2)
191200

192201
m_database.connect.assert_called_once_with(database='template1', host='127.0.0.1', password='bar', port='5432', user='foo')

tox.ini

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,13 @@ envlist =
1818
py310-dj40-postgres
1919
py310-dj41-postgres
2020
py310-dj42-postgres
21+
py310-dj42-postgres3
2122
py310-dj32-mysql
2223
py310-dj40-mysql
2324
py310-dj41-mysql
2425
py310-dj42-mysql
2526
py310-djmaster-postgres
27+
py310-djmaster-postgres3
2628

2729
[testenv]
2830
commands = make test
@@ -38,6 +40,8 @@ passenv =
3840
setenv =
3941
postgres: DJANGO_EXTENSIONS_DATABASE_ENGINE = {env:DJANGO_EXTENSIONS_DATABASE_ENGINE:django.db.backends.postgresql}
4042
postgres: DJANGO_EXTENSIONS_DATABASE_NAME = {env:DJANGO_EXTENSIONS_DATABASE_NAME:django_extensions_test}
43+
postgres3: DJANGO_EXTENSIONS_DATABASE_ENGINE = {env:DJANGO_EXTENSIONS_DATABASE_ENGINE:django.db.backends.postgresql}
44+
postgres3: DJANGO_EXTENSIONS_DATABASE_NAME = {env:DJANGO_EXTENSIONS_DATABASE_NAME:django_extensions_test}
4145
mysql: DJANGO_EXTENSIONS_DATABASE_ENGINE = {env:DJANGO_EXTENSIONS_DATABASE_ENGINE:django.db.backends.mysql}
4246
mysql: DJANGO_EXTENSIONS_DATABASE_NAME = {env:DJANGO_EXTENSIONS_DATABASE_NAME:django_extensions_test}
4347

@@ -50,6 +54,7 @@ deps =
5054
dj42: Django>=4.2,<5.0
5155
djmaster: https://github.com/django/django/archive/refs/heads/main.zip
5256
postgres: psycopg2-binary
57+
postgres3: psycopg[binary,pool]
5358
mysql: mysqlclient
5459

5560
[testenv:precommit]

0 commit comments

Comments
 (0)
0