Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Better compat with Model and Django Migrations #31

Open
wants to merge 16 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ class Address(CompositeType):
country = models.CharField(max_length=50)

class Meta:
db_type = 'x_address' # Required
db_table = 'x_address' # Required


class Person(models.Model):
Expand Down Expand Up @@ -76,7 +76,7 @@ class Card(CompositeType):
rank = models.CharField(max_length=2)

class Meta:
db_type = 'card'
db_table = 'card'


class Hand(models.Model):
Expand All @@ -94,13 +94,13 @@ class Point(CompositeType):
y = models.IntegerField()

class Meta:
db_type = 'x_point' # Postgres already has a point type
db_table = 'x_point' # Postgres already has a point type


class Box(CompositeType):
"""An axis-aligned box on the cartesian plane."""
class Meta:
db_type = 'x_box' # Postgres already has a box type
db_table = 'x_box' # Postgres already has a box type

top_left = Point.Field()
bottom_right = Point.Field()
Expand Down
9 changes: 7 additions & 2 deletions postgres_composite_types/caster.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,10 @@
from typing import TYPE_CHECKING, Type

from psycopg2.extras import CompositeCaster

if TYPE_CHECKING:
from .composite_type import CompositeType

__all__ = ["BaseCaster"]


Expand All @@ -9,7 +14,7 @@ class BaseCaster(CompositeCaster):
instance.
"""

Meta = None
_composite_type_model: Type["CompositeType"]

def make(self, values):
return self.Meta.model(*values)
return self._composite_type_model(*values)
194 changes: 89 additions & 105 deletions postgres_composite_types/composite_type.py
Original file line number Diff line number Diff line change
@@ -1,33 +1,25 @@
import inspect
import logging
import sys
from typing import Type

from django.db import models
from django.db.backends.postgresql.base import (
DatabaseWrapper as PostgresDatabaseWrapper,
)
from django.db import connections, models
from django.db.backends.signals import connection_created
from django.db.models.base import ModelBase
from django.db.models.manager import EmptyManager
from django.db.models.signals import post_migrate
from psycopg2 import ProgrammingError
from psycopg2.extensions import ISQLQuote, register_adapter
from psycopg2.extras import CompositeCaster, register_composite
from psycopg2.extras import register_composite

from .caster import BaseCaster
from .fields import BaseField
from .operations import BaseOperation
from .fields import BaseField, DummyField
from .quoting import QuotedCompositeType

LOGGER = logging.getLogger(__name__)

__all__ = ["CompositeType"]


def _add_class_to_module(cls, module_name):
cls.__module__ = module_name
module = sys.modules[module_name]
setattr(module, cls.__name__, cls)


class CompositeTypeMeta(type):
class CompositeTypeMeta(ModelBase):
"""Metaclass for Type."""

@classmethod
Expand All @@ -40,6 +32,7 @@ def __prepare__(cls, name, bases):
"""
return {}

# pylint:disable=arguments-differ
def __new__(cls, name, bases, attrs):
# Only apply the metaclass to our subclasses
if name == "CompositeType":
Expand All @@ -52,94 +45,62 @@ def __new__(cls, name, bases, attrs):
raise TypeError("Composite types cannot contain " "related fields")

if isinstance(value, models.Field):
field = attrs.pop(field_name)
field = attrs[field_name]
field.set_attributes_from_name(field_name)
fields.append((field_name, field))

# retrieve the Meta from our declaration
try:
meta_obj = attrs.pop("Meta")
meta_obj = attrs["Meta"]
except KeyError as exc:
raise TypeError(f'{name} has no "Meta" class') from exc

try:
meta_obj.db_type
meta_obj.db_table
except AttributeError as exc:
raise TypeError(f"{name}.Meta.db_type is required.") from exc

meta_obj.fields = fields
raise TypeError(f"{name}.Meta.db_table is required.") from exc

# create the field for this Type
attrs["Field"] = type(f"{name}Field", (BaseField,), {"Meta": meta_obj})

# add field class to the module in which the composite type class lives
# this is required for migrations to work
_add_class_to_module(attrs["Field"], attrs["__module__"])

# create the database operation for this type
attrs["Operation"] = type(
f"Create{name}Type", (BaseOperation,), {"Meta": meta_obj}
)

# create the caster for this type
attrs["Caster"] = type(f"{name}Caster", (BaseCaster,), {"Meta": meta_obj})
attrs["Field"] = type(f"{name}.Field", (BaseField,), {})

new_cls = super().__new__(cls, name, bases, attrs)
new_cls._meta = meta_obj
attrs[DummyField.name] = DummyField(primary_key=True, serialize=False)

meta_obj.model = new_cls
# Use an EmptyManager for everything as types cannot be queried.
meta_obj.default_manager_name = "objects"
meta_obj.base_manager_name = "objects"
attrs["objects"] = EmptyManager(model=None) # type: ignore

return new_cls
ret = super().__new__(cls, name, bases, attrs)
ret.Field._composite_type_model = ret # type: ignore
return ret

def __init__(cls, name, bases, attrs):
super().__init__(name, bases, attrs)
if name == "CompositeType":
return

cls._capture_descriptors() # pylint:disable=no-value-for-parameter
# pylint:disable=no-value-for-parameter
cls._connect_signals()

# Register the type on the first database connection
connection_created.connect(
receiver=cls.database_connected, dispatch_uid=cls._meta.db_type
)

def _capture_descriptors(cls):
"""Work around for not being able to call contribute_to_class.

Too much code to fake in our meta objects etc to be able to call
contribute_to_class directly, but we still want fields to be able
to set custom type descriptors. So we fake a model instead, with the
same fields as the composite type, and extract any custom descriptors
on that.
def _on_signal_register_type(cls, signal, sender, connection=None, **kwargs):
"""
Attempt registering the type after a migration succeeds.
"""
from django.db.backends.postgresql.base import DatabaseWrapper

attrs = dict(cls._meta.fields)

# we need to build a unique app label and model name combination for
# every composite type so django doesn't complain about model reloads
class Meta:
app_label = cls.__module__

attrs["__module__"] = cls.__module__
attrs["Meta"] = Meta
model_name = f"_Fake{cls.__name__}Model"
if connection is None:
connection = connections["default"]

fake_model = type(model_name, (models.Model,), attrs)
for field_name, _ in cls._meta.fields:
attr = getattr(fake_model, field_name)
if inspect.isdatadescriptor(attr):
setattr(cls, field_name, attr)
if isinstance(connection, DatabaseWrapper):
# On-connect, register the QuotedCompositeType with psycopg2.
# This is what to do when the type is going in to the database
register_adapter(cls, QuotedCompositeType)

def database_connected(cls, signal, sender, connection, **kwargs):
"""
Register this type with the database the first time a connection is
made.
"""
if isinstance(connection, PostgresDatabaseWrapper):
# Try to register the type. If the type has not been created in a
# migration, the registration will fail. The type will be
# Now try to register the type. If the type has not been created
# in a migration, the registration will fail. The type will be
# registered as part of the migration, so hopefully the migration
# will run soon.

try:
cls.register_composite(connection)
except ProgrammingError as exc:
Expand All @@ -150,11 +111,35 @@ def database_connected(cls, signal, sender, connection, **kwargs):
cls.__name__,
exc,
)
else:
# Registration succeeded.Disconnect the signals now.
cls._disconnect_signals() # pylint:disable=no-value-for-parameter

def _connect_signals(cls):
type_id = cls._meta.db_table

# Disconnect the signal now - only need to register types on the
# initial connection
# Register the type on the first database connection
connection_created.connect(
receiver=cls._on_signal_register_type, dispatch_uid=f"connect:{type_id}"
)

# Also register on post-migrate.
# This ensures that, if the on-connect signal failed due to a migration
# not having run yet, running the migration will still register it,
# even if in the same session (this can happen in tests for example).
# dispatch_uid needs to be distinct from the one on connection_created.
post_migrate.connect(
receiver=cls._on_signal_register_type,
dispatch_uid=f"post_migrate:{type_id}",
)

def _disconnect_signals(cls):
type_id = cls._meta.db_table
connection_created.disconnect(
cls.database_connected, dispatch_uid=cls._meta.db_type
cls._on_signal_register_type, dispatch_uid=f"connect:{type_id}"
)
post_migrate.disconnect(
cls._on_signal_register_type, dispatch_uid=f"post_migrate:{type_id}"
)


Expand All @@ -163,22 +148,21 @@ class CompositeType(metaclass=CompositeTypeMeta):
A new composite type stored in Postgres.
"""

_meta = None

# The database connection this type is registered with
registered_connection = None
_meta: Type

def __init__(self, *args, **kwargs):
if args and kwargs:
raise RuntimeError("Specify either args or kwargs but not both.")

# Initialise blank values for anyone expecting them
for name, _ in self._meta.fields:
setattr(self, name, None)
fields = self.get_fields()
for field in fields:
setattr(self, field.name, None)

# Unpack any args as if they came from the type
for (name, _), arg in zip(self._meta.fields, args):
setattr(self, name, arg)
for field, arg in zip(fields, args):
setattr(self, field.name, arg)

for name, value in kwargs.items():
setattr(self, name, value)
Expand All @@ -189,23 +173,23 @@ def __repr__(self):

def __to_tuple__(self):
return tuple(
field.get_prep_value(getattr(self, name))
for name, field in self._meta.fields
field.get_prep_value(getattr(self, field.name))
for field in self.get_fields()
)

def __to_dict__(self):
return {
name: field.get_prep_value(getattr(self, name))
for name, field in self._meta.fields
field.name: field.get_prep_value(getattr(self, field.name))
for field in self.get_fields()
}

def __eq__(self, other):
if not isinstance(other, CompositeType):
return False
if self._meta.model != other._meta.model:
return False
for name, _ in self._meta.fields:
if getattr(self, name) != getattr(other, name):
for field in self.get_fields():
if getattr(self, field.name) != getattr(other, field.name):
return False
return True

Expand All @@ -227,11 +211,10 @@ def register_composite(cls, connection):

with connection.temporary_connection() as cur:
# This is what to do when the type is coming out of the database
register_composite(
cls._meta.db_type, cur, globally=True, factory=cls.Caster
)
# This is what to do when the type is going in to the database
register_adapter(cls, QuotedCompositeType)
# We create a custom class subclassing BaseCaster (see caster.py),
# and set _composite_type_model attribute accordingly.
caster = type("Caster", (BaseCaster,), {"_composite_type_model": cls})
register_composite(cls._meta.db_table, cur, globally=True, factory=caster)

def __conform__(self, protocol):
"""
Expand All @@ -251,12 +234,13 @@ class Field(BaseField):
Placeholder for the field that will be produced for this type.
"""

class Operation(BaseOperation):
"""
Placeholder for the DB operation that will be produced for this type.
"""
# pylint:disable=invalid-name
def _get_next_or_previous_by_FIELD(self):
pass

class Caster(CompositeCaster):
"""
Placeholder for the caster that will be produced for this type
"""
@classmethod
def check(cls, **kwargs):
return []

def get_fields(self):
return self._meta.fields
Loading