Skip to content

Commit 85eb1e9

Browse files
timgrahamtimgraham
committed
add support for psycopg3
1 parent e26e81e commit 85eb1e9

File tree

12 files changed

+133
-22
lines changed

12 files changed

+133
-22
lines changed

‎README.md

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,12 @@
22

33
## Prerequisites
44

5-
You must install either:
5+
You must install:
6+
7+
* [psycopg](https://pypi.org/project/psycopg/), which may have some
8+
prerequisites [depending on which version you use](https://www.psycopg.org/psycopg3/docs/basic/install.html).
9+
10+
You can also use either:
611

712
* [psycopg2](https://pypi.org/project/psycopg2/), which has some
813
[prerequisites](https://www.psycopg.org/docs/install.html#prerequisites) of

‎django_cockroachdb/base.py

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -7,14 +7,12 @@
77
from django.utils.functional import cached_property
88

99
try:
10-
import psycopg2 # noqa
11-
import psycopg2.extensions # noqa
12-
import psycopg2.extras # noqa
13-
except ImportError as err:
14-
raise ImproperlyConfigured(
15-
'Error loading psycopg2 module.\n'
16-
'Did you install psycopg2 or psycopg2-binary?'
17-
) from err
10+
try:
11+
import psycopg # noqa
12+
except ImportError:
13+
import psycopg2 # noqa
14+
except ImportError:
15+
raise ImproperlyConfigured("Error loading psycopg or psycopg2 module")
1816

1917
from django.db.backends.postgresql.base import (
2018
DatabaseWrapper as PostgresDatabaseWrapper,

‎django_cockroachdb/features.py

Lines changed: 58 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -117,10 +117,6 @@ def django_test_expected_failures(self):
117117
'serializers.test_data.SerializerDataTests.test_yaml_serializer',
118118
# No sequence for AutoField in CockroachDB.
119119
'introspection.tests.IntrospectionTests.test_sequence_list',
120-
# Unsupported query: unsupported binary operator: <int> / <int>:
121-
# https://github.com/cockroachdb/django-cockroachdb/issues/21
122-
'expressions.tests.ExpressionOperatorTests.test_lefthand_division',
123-
'expressions.tests.ExpressionOperatorTests.test_right_hand_division',
124120
# CockroachDB doesn't support disabling constraints:
125121
# https://github.com/cockroachdb/cockroach/issues/19444
126122
'auth_tests.test_views.UUIDUserTests.test_admin_password_change',
@@ -196,6 +192,63 @@ def django_test_expected_failures(self):
196192
# https://github.com/cockroachdb/cockroach/issues/73587#issuecomment-988408190
197193
'aggregation.tests.AggregateTestCase.test_aggregation_default_using_decimal_from_database',
198194
})
195+
if self.uses_server_side_binding:
196+
expected_failures.update({
197+
# could not determine data type of placeholder:
198+
# https://github.com/cockroachdb/cockroach/issues/91396
199+
'backends.tests.EscapingChecks.test_parameter_escaping',
200+
'backends.tests.EscapingChecksDebug.test_parameter_escaping',
201+
'expressions.tests.BasicExpressionsTests.test_annotate_values_filter',
202+
'expressions_case.tests.CaseDocumentationExamples.test_lookup_example',
203+
'expressions_case.tests.CaseDocumentationExamples.test_simple_example',
204+
'expressions_case.tests.CaseExpressionTests.test_aggregation_empty_cases',
205+
'expressions_case.tests.CaseExpressionTests.test_annotate',
206+
'expressions_case.tests.CaseExpressionTests.test_annotate_exclude',
207+
'expressions_case.tests.CaseExpressionTests.test_annotate_values_not_in_order_by',
208+
'expressions_case.tests.CaseExpressionTests.test_annotate_with_aggregation_in_condition',
209+
'expressions_case.tests.CaseExpressionTests.test_annotate_with_aggregation_in_predicate',
210+
'expressions_case.tests.CaseExpressionTests.test_annotate_with_annotation_in_condition',
211+
'expressions_case.tests.CaseExpressionTests.test_annotate_with_annotation_in_predicate',
212+
'expressions_case.tests.CaseExpressionTests.test_annotate_with_empty_when',
213+
'expressions_case.tests.CaseExpressionTests.test_annotate_with_expression_as_condition',
214+
'expressions_case.tests.CaseExpressionTests.test_annotate_with_full_when',
215+
'expressions_case.tests.CaseExpressionTests.test_annotate_with_join_in_condition',
216+
'expressions_case.tests.CaseExpressionTests.test_annotate_with_join_in_predicate',
217+
'expressions_case.tests.CaseExpressionTests.test_case_reuse',
218+
'expressions_case.tests.CaseExpressionTests.test_combined_q_object',
219+
'expressions_case.tests.CaseExpressionTests.test_lookup_different_fields',
220+
'expressions_case.tests.CaseExpressionTests.test_lookup_in_condition',
221+
'expressions_case.tests.CaseExpressionTests.test_update_generic_ip_address',
222+
'lookup.tests.LookupQueryingTests.test_conditional_expression',
223+
'ordering.tests.OrderingTests.test_order_by_constant_value',
224+
'queries.test_bulk_update.BulkUpdateNoteTests.test_batch_size',
225+
'queries.test_bulk_update.BulkUpdateNoteTests.test_multiple_fields',
226+
'queries.test_bulk_update.BulkUpdateNoteTests.test_simple',
227+
'queries.test_bulk_update.BulkUpdateTests.test_custom_pk',
228+
'queries.test_bulk_update.BulkUpdateTests.test_database_routing',
229+
'queries.test_bulk_update.BulkUpdateTests.test_database_routing_batch_atomicity',
230+
'queries.test_bulk_update.BulkUpdateTests.test_falsey_pk_value',
231+
'queries.test_bulk_update.BulkUpdateTests.test_inherited_fields',
232+
'queries.test_bulk_update.BulkUpdateTests.test_large_batch',
233+
'queries.test_bulk_update.BulkUpdateTests.test_updated_rows_when_passing_duplicates',
234+
'queries.test_q.QCheckTests.test_expression',
235+
# unsupported binary operator: <interval> / <decimal>
236+
'expressions.tests.FTimeDeltaTests.test_durationfield_multiply_divide',
237+
# InvalidParameterValue: unsupported binary operator: <int4> / <float>
238+
'queries.tests.Ticket23605Tests.test_ticket_23605',
239+
# InvalidParameterValue: unsupported binary operator: <int2> + <float>
240+
'annotations.tests.NonAggregateAnnotationTestCase.test_combined_annotation_commutative',
241+
# incompatible COALESCE expressions: unsupported binary
242+
# operator: <decimal> / <float> (desired <decimal>)
243+
'aggregation.tests.AggregateTestCase.test_aggregation_default_passed_another_aggregate',
244+
})
245+
else:
246+
expected_failures.update({
247+
# Unsupported query: unsupported binary operator: <int> / <int>:
248+
# https://github.com/cockroachdb/django-cockroachdb/issues/21
249+
'expressions.tests.ExpressionOperatorTests.test_lefthand_division',
250+
'expressions.tests.ExpressionOperatorTests.test_right_hand_division',
251+
})
199252
return expected_failures
200253

201254
@cached_property
@@ -224,7 +277,7 @@ def django_test_skips(self):
224277
# output in the logs:
225278
# Exception in thread Thread-1:
226279
# ...
227-
# psycopg2.errors.SerializationFailure: restart transaction:
280+
# psycopg.errors.SerializationFailure: restart transaction:
228281
# TransactionRetryWithProtoRefreshError: WriteTooOldError: write
229282
# at timestamp 1598314405.858850941,0 too old; wrote at
230283
# 1598314405.883337663,1

‎django_cockroachdb/operations.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
from django.db.backends.postgresql.operations import (
33
DatabaseOperations as PostgresDatabaseOperations,
44
)
5+
from django.db.backends.postgresql.psycopg_any import is_psycopg3
56

67

78
class DatabaseOperations(PostgresDatabaseOperations):
@@ -16,6 +17,19 @@ class DatabaseOperations(PostgresDatabaseOperations):
1617
'AutoField': (-9223372036854775808, 9223372036854775807),
1718
'BigAutoField': (-9223372036854775808, 9223372036854775807),
1819
}
20+
21+
if is_psycopg3:
22+
from psycopg.types import numeric
23+
24+
integerfield_type_map = {
25+
"SmallIntegerField": numeric.Int2,
26+
"IntegerField": numeric.Int8,
27+
"BigIntegerField": numeric.Int8,
28+
"PositiveSmallIntegerField": numeric.Int2,
29+
"PositiveIntegerField": numeric.Int8,
30+
"PositiveBigIntegerField": numeric.Int8,
31+
}
32+
1933
explain_options = frozenset(['DISTSQL', 'OPT', 'TYPES', 'VEC', 'VERBOSE'])
2034

2135
def deferrable_sql(self):
@@ -25,7 +39,7 @@ def deferrable_sql(self):
2539

2640
def adapt_datetimefield_value(self, value):
2741
"""
28-
Add a timezone to datetimes so that psycopg2 will cast it to
42+
Add a timezone to datetimes so that psycopg will cast it to
2943
TIMESTAMPTZ (as cockroach expects) rather than TIMESTAMP.
3044
"""
3145
# getattr() guards against F() objects which don't have tzinfo.

‎django_cockroachdb/schema.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,15 @@ class DatabaseSchemaEditor(PostgresDatabaseSchemaEditor):
2424
# statement. This isn't supported by CockroachDB.
2525
sql_update_with_default = "UPDATE %(table)s SET %(column)s = %(default)s WHERE %(column)s IS NULL"
2626

27+
def __enter__(self):
28+
super().__enter__()
29+
# As long as DatabaseFeatures.can_rollback_ddl = False, compose() may
30+
# fail if connection is None as per
31+
# https://github.com/django/django/pull/15687#discussion_r1038175823.
32+
# See also https://github.com/django/django/pull/15687#discussion_r1041503991.
33+
self.connection.ensure_connection()
34+
return self
35+
2736
def add_index(self, model, index, concurrently=False):
2837
if index.contains_expressions and not self.connection.features.supports_expression_indexes:
2938
return None

‎django_cockroachdb_gis/base.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,7 @@
1+
from django.contrib.gis.db.backends.postgis.base import (
2+
DatabaseWrapper as PostGISDatabaseWrapper,
3+
)
4+
15
from django_cockroachdb.base import DatabaseWrapper as CockroachDatabaseWrapper
26

37
from .features import DatabaseFeatures
@@ -6,8 +10,13 @@
610
from .schema import DatabaseSchemaEditor
711

812

9-
class DatabaseWrapper(CockroachDatabaseWrapper):
13+
class DatabaseWrapper(CockroachDatabaseWrapper, PostGISDatabaseWrapper):
1014
SchemaEditorClass = DatabaseSchemaEditor
1115
features_class = DatabaseFeatures
1216
introspection_class = DatabaseIntrospection
1317
ops_class = DatabaseOperations
18+
19+
def __init__(self, *args, **kwargs):
20+
# Skip PostGISDatabaseWrapper.__init__() to work around
21+
# https://code.djangoproject.com/ticket/34344.
22+
CockroachDatabaseWrapper.__init__(self, *args, **kwargs)

‎django_cockroachdb_gis/features.py

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -37,9 +37,6 @@ def django_test_expected_failures(self):
3737
'gis_tests.geoapp.tests.GeoLookupTest.test_relate_lookup',
3838
# NotSupportedError: this box2d comparison operator is experimental
3939
'gis_tests.geoapp.tests.GeoLookupTest.test_contains_contained_lookups',
40-
# unknown signature: st_dwithin(geography, geometry, decimal) (desired <bool>)
41-
# https://github.com/cockroachdb/cockroach/issues/53720
42-
'gis_tests.geogapp.tests.GeographyTest.test02_distance_lookup',
4340
# unknown signature: st_distancespheroid(geometry, geometry, string)
4441
# https://github.com/cockroachdb/cockroach/issues/48922#issuecomment-693096502
4542
'gis_tests.distapp.tests.DistanceTest.test_distance_lookups_with_expression_rhs',
@@ -74,4 +71,15 @@ def django_test_expected_failures(self):
7471
# https://github.com/cockroachdb/cockroach/issues/47420#issuecomment-969578772
7572
'gis_tests.gis_migrations.test_operations.OperationTests.test_add_3d_field_opclass',
7673
})
74+
if self.uses_server_side_binding:
75+
expected_failures.update({
76+
# unknown signature: st_scale(geometry, int, int)
77+
'gis_tests.geoapp.test_functions.GISFunctionsTests.test_scale',
78+
})
79+
else:
80+
expected_failures.update({
81+
# unknown signature: st_dwithin(geography, geometry, decimal) (desired <bool>)
82+
# https://github.com/cockroachdb/cockroach/issues/53720
83+
'gis_tests.geogapp.tests.GeographyTest.test02_distance_lookup',
84+
})
7785
return expected_failures

‎teamcity-build/build-teamcity-22.1.sh

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#!/usr/bin/env bash
22
set -x
33

4+
export USE_PSYCOPG2=1
45
./teamcity-build/build-teamcity.sh "v22.1.18"

‎teamcity-build/build-teamcity-22.2-psycopg2.sh

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#!/usr/bin/env bash
22
set -x
33

4-
# Not supported
4+
export USE_PSYCOPG2=1
5+
./teamcity-build/build-teamcity.sh "v22.2.7"

‎teamcity-build/build-teamcity-22.2-server-side-binding.sh

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#!/usr/bin/env bash
22
set -x
33

4-
# Not supported
4+
export USE_SERVER_SIDE_BINDING=1
5+
./teamcity-build/build-teamcity.sh "v22.2.7"

0 commit comments

Comments
 (0)