1
0
Fork 0
mirror of https://gitlab.com/bramw/baserow.git synced 2025-04-15 01:28:30 +00:00

psycopg3 tests

This commit is contained in:
Cezary Statkiewicz 2025-02-25 15:54:15 +01:00
parent 87b7f19d75
commit 57958fb44c
4 changed files with 186 additions and 173 deletions
backend
src/baserow
contrib/database
core
tests/baserow/contrib/database/field

View file

@ -4,9 +4,9 @@ from django.core.exceptions import FieldDoesNotExist
from django.db import ProgrammingError
from django.db.models.signals import post_migrate, pre_migrate
from baserow.contrib.database.fields.utils.pg_datetime import pg_init # noqa: F401
from baserow.contrib.database.table.cache import clear_generated_model_cache
from baserow.contrib.database.table.operations import RestoreDatabaseTableOperationType
from baserow.core.psycopg import is_psycopg3
from baserow.core.registries import (
application_type_registry,
object_scope_type_registry,
@ -972,7 +972,13 @@ class DatabaseConfig(AppConfig):
import baserow.contrib.database.table.receivers # noqa: F401
import baserow.contrib.database.views.tasks # noqa: F401
pg_init()
# date/datetime min/max year handling - we need that for psycopg 3.x only
if is_psycopg3:
from baserow.contrib.database.fields.utils.pg_datetime import ( # noqa: F401
pg_init,
)
pg_init()
# noinspection PyPep8Naming

View file

@ -1,113 +1,105 @@
from django.db.backends.signals import connection_created
from baserow.core.psycopg import is_psycopg3, psycopg
import psycopg
from psycopg.types.datetime import (
DataError,
DateBinaryLoader,
DateLoader,
TimestampBinaryLoader,
TimestampLoader,
TimestamptzBinaryLoader,
TimestamptzLoader,
)
if is_psycopg3:
from django.db.backends.signals import connection_created
from baserow.core.psycopg import (
DataError,
DateBinaryLoader,
DateLoader,
TimestampBinaryLoader,
TimestampLoader,
TimestamptzBinaryLoader,
TimestamptzLoader,
)
# sentinel
class DateOverflowPlaceholder:
INVALID_DATE = "infinity"
# sentinel
class DateOverflowPlaceholder:
INVALID_DATE = "infinity"
def isoformat(self):
return self.INVALID_DATE
def isoformat(self):
return self.INVALID_DATE
def for_json(self):
return self.INVALID_DATE
def for_json(self):
return self.INVALID_DATE
def __str__(self):
return self.INVALID_DATE
def __str__(self):
return self.INVALID_DATE
def __cmp__(self, other):
return isinstance(other, self.__class__) or self.INVALID_DATE == other
def __cmp__(self, other):
return isinstance(other, self.__class__) or self.INVALID_DATE == other
DATE_OVERFLOW = DateOverflowPlaceholder()
DATE_OVERFLOW = DateOverflowPlaceholder()
class _DateOverflowLoaderMixin:
def load(self, data):
try:
return super().load(data)
except DataError:
return DATE_OVERFLOW
class _TimestamptzOverflowLoaderMixin:
timezone = None
class _DateOverflowLoaderMixin:
def load(self, data):
try:
return super().load(data)
except DataError:
return DATE_OVERFLOW
def load(self, data):
try:
res = super().load(data)
return res.replace(tzinfo=self.timezone)
except DataError:
return DATE_OVERFLOW
class BaserowDateLoader(_DateOverflowLoaderMixin, DateLoader):
pass
class _TimestamptzOverflowLoaderMixin:
timezone = None
class BaserowDateBinaryLoader(_DateOverflowLoaderMixin, DateBinaryLoader):
pass
def load(self, data):
try:
res = super().load(data)
return res.replace(tzinfo=self.timezone)
except DataError:
return DATE_OVERFLOW
class BaserowTimestampLoader(_DateOverflowLoaderMixin, TimestampLoader):
pass
class BaserowTimestampBinaryLoader(_DateOverflowLoaderMixin, TimestampBinaryLoader):
pass
class BaserowDateLoader(_DateOverflowLoaderMixin, DateLoader):
pass
def pg_init():
"""
Registers loaders for psycopg3 to handle date overflow.
:return:
"""
class BaserowDateBinaryLoader(_DateOverflowLoaderMixin, DateBinaryLoader):
pass
psycopg.adapters.register_loader("date", BaserowDateLoader)
psycopg.adapters.register_loader("date", BaserowDateBinaryLoader)
psycopg.adapters.register_loader("timestamp", BaserowTimestampLoader)
psycopg.adapters.register_loader("timestamp", BaserowTimestampBinaryLoader)
class BaserowTimestampLoader(_DateOverflowLoaderMixin, TimestampLoader):
pass
# psycopg3 and timezones allow per-connection / per-cursor adapting. This is
# done in django/db/backends/postgresql/psycopg_any.py in a hook that
# registries tz aware adapter for each connection/cursor.
# We can re-register our loaders here, but note that this will work on
# per-connection tz setting. Cursors still will use django-provided adapters
def register_context(signal, sender, connection, **kwargs):
register_on_connection(connection)
connection_created.connect(register_context)
class BaserowTimestampBinaryLoader(_DateOverflowLoaderMixin, TimestampBinaryLoader):
pass
def register_on_connection(connection):
"""
Registers timestamptz pg type loaders for a connection.
:param connection:
:return:
"""
def pg_init():
"""
Registers loaders for psycopg3 to handle date overflow.
ctx = connection.connection.adapters
:return:
"""
class SpecificTzLoader(_TimestamptzOverflowLoaderMixin, TimestamptzLoader):
timezone = connection.timezone
psycopg.adapters.register_loader("date", BaserowDateLoader)
psycopg.adapters.register_loader("date", BaserowDateBinaryLoader)
class SpecificTzBinaryLoader(
_TimestamptzOverflowLoaderMixin, TimestamptzBinaryLoader
):
timezone = connection.timezone
psycopg.adapters.register_loader("timestamp", BaserowTimestampLoader)
psycopg.adapters.register_loader("timestamp", BaserowTimestampBinaryLoader)
# psycopg3 and timezones allow per-connection / per-cursor adapting. This is done in
# django/db/backends/postgresql/psycopg_any.py in a hook that registries tz aware
# adapter for each connection/cursor.
# We can re-register our loaders here, but note that this will work on
# per-connection tz setting. Cursors still will use django-provided adapters
def register_context(signal, sender, connection, **kwargs):
register_on_connection(connection)
connection_created.connect(register_context)
def register_on_connection(connection):
"""
Registers timestamptz pg type loaders for a connection.
:param connection:
:return:
"""
ctx = connection.connection.adapters
class SpecificTzLoader(_TimestamptzOverflowLoaderMixin, TimestamptzLoader):
timezone = connection.timezone
class SpecificTzBinaryLoader(
_TimestamptzOverflowLoaderMixin, TimestamptzBinaryLoader
):
timezone = connection.timezone
ctx.adapters.register_loader("timestamptz", SpecificTzLoader)
ctx.adapters.register_loader("timestamptz", SpecificTzBinaryLoader)
ctx.adapters.register_loader("timestamptz", SpecificTzLoader)
ctx.adapters.register_loader("timestamptz", SpecificTzBinaryLoader)

View file

@ -3,6 +3,18 @@ from django.db.backends.postgresql.psycopg_any import is_psycopg3
if is_psycopg3:
import psycopg # noqa: F401
from psycopg import sql # noqa: F401
# used for date type mapping
from psycopg.types.datetime import ( # noqa: F401
DataError,
DateBinaryLoader,
DateLoader,
TimestampBinaryLoader,
TimestampLoader,
TimestamptzBinaryLoader,
TimestamptzLoader,
)
else:
import psycopg2 as psycopg # noqa: F401
from psycopg2 import sql # noqa: F401

View file

@ -11,12 +11,9 @@ from baserow.contrib.database.fields.handler import FieldHandler
from baserow.contrib.database.fields.models import DateField, TextField
from baserow.contrib.database.fields.registries import field_type_registry
from baserow.contrib.database.fields.utils import DeferredForeignKeyUpdater
from baserow.contrib.database.fields.utils.pg_datetime import (
DATE_OVERFLOW,
register_on_connection,
)
from baserow.contrib.database.rows.handler import RowHandler
from baserow.contrib.database.views.handler import ViewHandler
from baserow.core.psycopg import is_psycopg3
from baserow.core.registries import ImportExportConfig
@ -744,98 +741,104 @@ def test_get_group_by_metadata_in_rows_with_date_field(data_fixture):
}
@pytest.mark.django_db
def test_date_field_overflow(settings, data_fixture):
user = data_fixture.create_user()
table = data_fixture.create_database_table(user=user)
field_handler = FieldHandler()
row_handler = RowHandler()
date_field = field_handler.create_field(
user=user,
table=table,
type_name="text",
name="Date",
)
invalid_date_value = "19999-01-01"
row = row_handler.create_row(
user=user, table=table, values={date_field.db_column: invalid_date_value}
)
assert getattr(row, date_field.db_column, None) == invalid_date_value
date_field = field_handler.update_field(
user=user, field=date_field, new_type_name="date", date_format="ISO"
if is_psycopg3:
from baserow.contrib.database.fields.utils.pg_datetime import (
DATE_OVERFLOW,
register_on_connection,
)
assert isinstance(
table.get_model().get_field_object(date_field.db_column)["field"], DateField
)
out = row_handler.get_rows(table.get_model(), [row.id])
assert len(out) == 1
assert getattr(out[0], date_field.db_column, None) is DATE_OVERFLOW
@pytest.mark.django_db
def test_date_field_overflow(settings, data_fixture):
user = data_fixture.create_user()
table = data_fixture.create_database_table(user=user)
date_field = field_handler.update_field(
user=user, field=date_field, new_type_name="text", date_format="ISO"
)
field_handler = FieldHandler()
row_handler = RowHandler()
table.refresh_from_db()
assert isinstance(
table.get_model().get_field_object(date_field.db_column)["field"], TextField
)
out = row_handler.get_rows(table.get_model(), [row.id])
assert len(out) == 1
assert getattr(out[0], date_field.db_column, None) == invalid_date_value
date_field = field_handler.create_field(
user=user,
table=table,
type_name="text",
name="Date",
)
invalid_date_value = "19999-01-01"
row = row_handler.create_row(
user=user, table=table, values={date_field.db_column: invalid_date_value}
)
assert getattr(row, date_field.db_column, None) == invalid_date_value
date_field = field_handler.update_field(
user=user, field=date_field, new_type_name="date", date_format="ISO"
)
@pytest.mark.django_db
def test_datetime_field_overflow(on_db_connection, data_fixture):
# manually register adapters, as signal-based registration will be called too late
on_db_connection(register_on_connection)
assert isinstance(
table.get_model().get_field_object(date_field.db_column)["field"], DateField
)
out = row_handler.get_rows(table.get_model(), [row.id])
assert len(out) == 1
assert getattr(out[0], date_field.db_column, None) is DATE_OVERFLOW
user = data_fixture.create_user()
table = data_fixture.create_database_table(user=user)
date_field = field_handler.update_field(
user=user, field=date_field, new_type_name="text", date_format="ISO"
)
field_handler = FieldHandler()
row_handler = RowHandler()
table.refresh_from_db()
assert isinstance(
table.get_model().get_field_object(date_field.db_column)["field"], TextField
)
out = row_handler.get_rows(table.get_model(), [row.id])
assert len(out) == 1
assert getattr(out[0], date_field.db_column, None) == invalid_date_value
date_field = field_handler.create_field(
user=user,
table=table,
type_name="text",
name="Date",
)
invalid_date_value = "19999-01-01 01:01"
row = row_handler.create_row(
user=user, table=table, values={date_field.db_column: invalid_date_value}
)
assert getattr(row, date_field.db_column, None) == invalid_date_value
@pytest.mark.django_db
def test_datetime_field_overflow(on_db_connection, data_fixture):
# manually register adapters, as signal-based registration will be called
# too late
on_db_connection(register_on_connection)
date_field = field_handler.update_field(
user=user,
field=date_field,
new_type_name="date",
date_format="ISO",
date_include_time=True,
date_time_format="24",
)
assert isinstance(
table.get_model().get_field_object(date_field.db_column)["field"], DateField
)
out = row_handler.get_rows(table.get_model(), [row.id])
assert len(out) == 1
user = data_fixture.create_user()
table = data_fixture.create_database_table(user=user)
assert getattr(out[0], date_field.db_column, None) is DATE_OVERFLOW
field_handler = FieldHandler()
row_handler = RowHandler()
date_field = field_handler.update_field(
user=user, field=date_field, new_type_name="text", date_format="ISO"
)
date_field = field_handler.create_field(
user=user,
table=table,
type_name="text",
name="Date",
)
invalid_date_value = "19999-01-01 01:01"
row = row_handler.create_row(
user=user, table=table, values={date_field.db_column: invalid_date_value}
)
assert getattr(row, date_field.db_column, None) == invalid_date_value
table.refresh_from_db()
assert isinstance(
table.get_model().get_field_object(date_field.db_column)["field"], TextField
)
out = row_handler.get_rows(table.get_model(), [row.id])
assert len(out) == 1
date_field = field_handler.update_field(
user=user,
field=date_field,
new_type_name="date",
date_format="ISO",
date_include_time=True,
date_time_format="24",
)
assert isinstance(
table.get_model().get_field_object(date_field.db_column)["field"], DateField
)
out = row_handler.get_rows(table.get_model(), [row.id])
assert len(out) == 1
assert getattr(out[0], date_field.db_column, None) == invalid_date_value
assert getattr(out[0], date_field.db_column, None) is DATE_OVERFLOW
date_field = field_handler.update_field(
user=user, field=date_field, new_type_name="text", date_format="ISO"
)
table.refresh_from_db()
assert isinstance(
table.get_model().get_field_object(date_field.db_column)["field"], TextField
)
out = row_handler.get_rows(table.get_model(), [row.id])
assert len(out) == 1
assert getattr(out[0], date_field.db_column, None) == invalid_date_value