Add CompositeType
This commit is contained in:
@@ -8,7 +8,7 @@ advantage of these datatypes you should use automatic data coercion. See :func:`
|
||||
|
||||
|
||||
ArrowType
|
||||
^^^^^^^^^
|
||||
---------
|
||||
|
||||
.. module:: sqlalchemy_utils.types.arrow
|
||||
|
||||
@@ -16,7 +16,7 @@ ArrowType
|
||||
|
||||
|
||||
ChoiceType
|
||||
^^^^^^^^^^
|
||||
----------
|
||||
|
||||
.. module:: sqlalchemy_utils.types.choice
|
||||
|
||||
@@ -24,15 +24,23 @@ ChoiceType
|
||||
|
||||
|
||||
ColorType
|
||||
^^^^^^^^^
|
||||
---------
|
||||
|
||||
.. module:: sqlalchemy_utils.types.color
|
||||
|
||||
.. autoclass:: ColorType
|
||||
|
||||
|
||||
CompositeType
|
||||
-------------
|
||||
|
||||
.. automodule:: sqlalchemy_utils.types.pg_composite
|
||||
|
||||
.. autoclass:: CompositeType
|
||||
|
||||
|
||||
CountryType
|
||||
^^^^^^^^^^^
|
||||
-----------
|
||||
|
||||
.. module:: sqlalchemy_utils.types.country
|
||||
|
||||
@@ -44,7 +52,7 @@ CountryType
|
||||
|
||||
|
||||
CurrencyType
|
||||
^^^^^^^^^^^^
|
||||
------------
|
||||
|
||||
.. module:: sqlalchemy_utils.types.currency
|
||||
|
||||
@@ -56,14 +64,14 @@ CurrencyType
|
||||
|
||||
|
||||
EncryptedType
|
||||
^^^^^^^^^^^^^
|
||||
-------------
|
||||
|
||||
.. module:: sqlalchemy_utils.types.encrypted
|
||||
|
||||
.. autoclass:: EncryptedType
|
||||
|
||||
JSONType
|
||||
^^^^^^^^
|
||||
--------
|
||||
|
||||
.. module:: sqlalchemy_utils.types.json
|
||||
|
||||
@@ -71,7 +79,7 @@ JSONType
|
||||
|
||||
|
||||
LocaleType
|
||||
^^^^^^^^^^
|
||||
----------
|
||||
|
||||
|
||||
.. module:: sqlalchemy_utils.types.locale
|
||||
@@ -80,7 +88,7 @@ LocaleType
|
||||
|
||||
|
||||
IPAddressType
|
||||
^^^^^^^^^^^^^
|
||||
-------------
|
||||
|
||||
.. module:: sqlalchemy_utils.types.ip_address
|
||||
|
||||
@@ -88,7 +96,7 @@ IPAddressType
|
||||
|
||||
|
||||
PasswordType
|
||||
^^^^^^^^^^^^
|
||||
------------
|
||||
|
||||
.. module:: sqlalchemy_utils.types.password
|
||||
|
||||
@@ -96,7 +104,7 @@ PasswordType
|
||||
|
||||
|
||||
PhoneNumberType
|
||||
^^^^^^^^^^^^^^^
|
||||
---------------
|
||||
|
||||
.. module:: sqlalchemy_utils.types.phone_number
|
||||
|
||||
@@ -104,7 +112,7 @@ PhoneNumberType
|
||||
|
||||
|
||||
ScalarListType
|
||||
^^^^^^^^^^^^^^
|
||||
--------------
|
||||
|
||||
.. module:: sqlalchemy_utils.types.scalar_list
|
||||
|
||||
@@ -112,7 +120,7 @@ ScalarListType
|
||||
|
||||
|
||||
TimezoneType
|
||||
^^^^^^^^^^^^
|
||||
------------
|
||||
|
||||
|
||||
.. module:: sqlalchemy_utils.types.timezone
|
||||
@@ -121,7 +129,7 @@ TimezoneType
|
||||
|
||||
|
||||
TSVectorType
|
||||
^^^^^^^^^^^^
|
||||
------------
|
||||
|
||||
.. module:: sqlalchemy_utils.types.ts_vector
|
||||
|
||||
@@ -129,7 +137,7 @@ TSVectorType
|
||||
|
||||
|
||||
URLType
|
||||
^^^^^^^
|
||||
-------
|
||||
|
||||
.. module:: sqlalchemy_utils.types.url
|
||||
|
||||
@@ -137,7 +145,7 @@ URLType
|
||||
|
||||
|
||||
UUIDType
|
||||
^^^^^^^^
|
||||
--------
|
||||
|
||||
|
||||
.. module:: sqlalchemy_utils.types.uuid
|
||||
@@ -147,7 +155,7 @@ UUIDType
|
||||
|
||||
|
||||
WeekDaysType
|
||||
^^^^^^^^^^^^
|
||||
------------
|
||||
|
||||
.. module:: sqlalchemy_utils.types.weekdays
|
||||
|
||||
|
@@ -6,54 +6,54 @@ Database helpers
|
||||
|
||||
|
||||
analyze
|
||||
^^^^^^^
|
||||
-------
|
||||
|
||||
.. autofunction:: analyze
|
||||
|
||||
|
||||
database_exists
|
||||
^^^^^^^^^^^^^^^
|
||||
---------------
|
||||
|
||||
.. autofunction:: database_exists
|
||||
|
||||
|
||||
create_database
|
||||
^^^^^^^^^^^^^^^
|
||||
---------------
|
||||
|
||||
.. autofunction:: create_database
|
||||
|
||||
|
||||
drop_database
|
||||
^^^^^^^^^^^^^
|
||||
-------------
|
||||
|
||||
.. autofunction:: drop_database
|
||||
|
||||
|
||||
has_index
|
||||
^^^^^^^^^
|
||||
---------
|
||||
|
||||
.. autofunction:: has_index
|
||||
|
||||
|
||||
has_unique_index
|
||||
^^^^^^^^^^^^^^^^
|
||||
----------------
|
||||
|
||||
.. autofunction:: has_unique_index
|
||||
|
||||
|
||||
json_sql
|
||||
^^^^^^^^
|
||||
--------
|
||||
|
||||
.. autofunction:: json_sql
|
||||
|
||||
|
||||
render_expression
|
||||
^^^^^^^^^^^^^^^^^
|
||||
-----------------
|
||||
|
||||
.. autofunction:: render_expression
|
||||
|
||||
|
||||
render_statement
|
||||
^^^^^^^^^^^^^^^^
|
||||
----------------
|
||||
|
||||
.. autofunction:: render_statement
|
||||
|
@@ -5,36 +5,36 @@ Foreign key helpers
|
||||
|
||||
|
||||
dependent_objects
|
||||
^^^^^^^^^^^^^^^^^
|
||||
-----------------
|
||||
|
||||
.. autofunction:: dependent_objects
|
||||
|
||||
|
||||
get_referencing_foreign_keys
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
----------------------------
|
||||
|
||||
.. autofunction:: get_referencing_foreign_keys
|
||||
|
||||
|
||||
group_foreign_keys
|
||||
^^^^^^^^^^^^^^^^^^
|
||||
------------------
|
||||
|
||||
.. autofunction:: group_foreign_keys
|
||||
|
||||
|
||||
is_indexed_foreign_key
|
||||
^^^^^^^^^^^^^^^^^^^^^^
|
||||
----------------------
|
||||
|
||||
.. autofunction:: is_indexed_foreign_key
|
||||
|
||||
|
||||
merge_references
|
||||
^^^^^^^^^^^^^^^^
|
||||
----------------
|
||||
|
||||
.. autofunction:: merge_references
|
||||
|
||||
|
||||
non_indexed_foreign_keys
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
------------------------
|
||||
|
||||
.. autofunction:: non_indexed_foreign_keys
|
||||
|
@@ -49,7 +49,7 @@ Generic relationship is a form of relationship that supports creating a 1 to man
|
||||
|
||||
|
||||
Inheritance
|
||||
^^^^^^^^^^^
|
||||
-----------
|
||||
|
||||
::
|
||||
|
||||
@@ -108,7 +108,7 @@ We can even test super types::
|
||||
|
||||
|
||||
Abstract base classes
|
||||
^^^^^^^^^^^^^^^^^^^^^
|
||||
---------------------
|
||||
|
||||
Generic relationships also allows using string arguments. When using generic_relationship with abstract base classes you need to set up the relationship using declared_attr decorator and string arguments.
|
||||
|
||||
@@ -140,7 +140,7 @@ Generic relationships also allows using string arguments. When using generic_rel
|
||||
|
||||
|
||||
Composite keys
|
||||
^^^^^^^^^^^^^^
|
||||
--------------
|
||||
|
||||
For some very rare cases you may need to use generic_relationships with composite primary keys. There is a limitation here though: you can only set up generic_relationship for similar composite primary key types. In other words you can't mix generic relationship to both composite keyed objects and single keyed objects.
|
||||
|
||||
|
@@ -3,7 +3,7 @@ Model mixins
|
||||
|
||||
|
||||
Timestamp
|
||||
^^^^^^^^^
|
||||
---------
|
||||
|
||||
.. module:: sqlalchemy_utils.models
|
||||
|
||||
|
@@ -5,108 +5,108 @@ ORM helpers
|
||||
|
||||
|
||||
escape_like
|
||||
^^^^^^^^^^^
|
||||
-----------
|
||||
|
||||
.. autofunction:: escape_like
|
||||
|
||||
|
||||
get_bind
|
||||
^^^^^^^^
|
||||
--------
|
||||
|
||||
.. autofunction:: get_bind
|
||||
|
||||
|
||||
get_class_by_table
|
||||
^^^^^^^^^^^^^^^^^^
|
||||
------------------
|
||||
|
||||
.. autofunction:: get_class_by_table
|
||||
|
||||
|
||||
get_column_key
|
||||
^^^^^^^^^^^^^^
|
||||
--------------
|
||||
|
||||
.. autofunction:: get_column_key
|
||||
|
||||
|
||||
get_columns
|
||||
^^^^^^^^^^^
|
||||
-----------
|
||||
|
||||
.. autofunction:: get_columns
|
||||
|
||||
|
||||
get_declarative_base
|
||||
^^^^^^^^^^^^^^^^^^^^
|
||||
--------------------
|
||||
|
||||
.. autofunction:: get_declarative_base
|
||||
|
||||
|
||||
get_hybrid_properties
|
||||
^^^^^^^^^^^^^^^^^^^^^
|
||||
---------------------
|
||||
|
||||
.. autofunction:: get_hybrid_properties
|
||||
|
||||
|
||||
get_mapper
|
||||
^^^^^^^^^^
|
||||
----------
|
||||
|
||||
.. autofunction:: get_mapper
|
||||
|
||||
|
||||
get_query_entities
|
||||
^^^^^^^^^^^^^^^^^^
|
||||
------------------
|
||||
|
||||
.. autofunction:: get_query_entities
|
||||
|
||||
|
||||
get_primary_keys
|
||||
^^^^^^^^^^^^^^^^
|
||||
----------------
|
||||
|
||||
.. autofunction:: get_primary_keys
|
||||
|
||||
|
||||
get_tables
|
||||
^^^^^^^^^^
|
||||
----------
|
||||
|
||||
.. autofunction:: get_tables
|
||||
|
||||
|
||||
has_changes
|
||||
^^^^^^^^^^^
|
||||
-----------
|
||||
|
||||
.. autofunction:: has_changes
|
||||
|
||||
|
||||
identity
|
||||
^^^^^^^^
|
||||
--------
|
||||
|
||||
.. autofunction:: identity
|
||||
|
||||
|
||||
is_loaded
|
||||
^^^^^^^^^
|
||||
---------
|
||||
|
||||
.. autofunction:: is_loaded
|
||||
|
||||
|
||||
make_order_by_deterministic
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
---------------------------
|
||||
|
||||
.. autofunction:: make_order_by_deterministic
|
||||
|
||||
|
||||
naturally_equivalent
|
||||
^^^^^^^^^^^^^^^^^^^^
|
||||
--------------------
|
||||
|
||||
.. autofunction:: naturally_equivalent
|
||||
|
||||
|
||||
quote
|
||||
^^^^^
|
||||
-----
|
||||
|
||||
.. autofunction:: quote
|
||||
|
||||
|
||||
sort_query
|
||||
^^^^^^^^^^
|
||||
----------
|
||||
|
||||
.. autofunction:: sort_query
|
||||
|
@@ -7,31 +7,31 @@ Range data types
|
||||
|
||||
|
||||
DateRangeType
|
||||
^^^^^^^^^^^^^
|
||||
-------------
|
||||
|
||||
.. autoclass:: DateRangeType
|
||||
|
||||
|
||||
DateTimeRangeType
|
||||
^^^^^^^^^^^^^^^^^
|
||||
-----------------
|
||||
|
||||
.. autoclass:: DateTimeRangeType
|
||||
|
||||
|
||||
IntRangeType
|
||||
^^^^^^^^^^^^
|
||||
------------
|
||||
|
||||
.. autoclass:: IntRangeType
|
||||
|
||||
|
||||
NumericRangeType
|
||||
^^^^^^^^^^^^^^^^
|
||||
----------------
|
||||
|
||||
.. autoclass:: NumericRangeType
|
||||
|
||||
|
||||
RangeComparator
|
||||
^^^^^^^^^^^^^^^
|
||||
---------------
|
||||
|
||||
.. autoclass:: RangeComparator
|
||||
:members:
|
||||
|
@@ -7,7 +7,7 @@ QueryChain
|
||||
.. automodule:: sqlalchemy_utils.query_chain
|
||||
|
||||
API
|
||||
^^^
|
||||
---
|
||||
|
||||
.. autoclass:: QueryChain
|
||||
:members:
|
||||
|
@@ -60,6 +60,8 @@ from .types import ( # noqa
|
||||
Choice,
|
||||
ChoiceType,
|
||||
ColorType,
|
||||
CompositeArray,
|
||||
CompositeType,
|
||||
CountryType,
|
||||
CurrencyType,
|
||||
DateRangeType,
|
||||
@@ -77,6 +79,8 @@ from .types import ( # noqa
|
||||
PasswordType,
|
||||
PhoneNumber,
|
||||
PhoneNumberType,
|
||||
register_composites,
|
||||
remove_composite_listeners,
|
||||
ScalarListException,
|
||||
ScalarListType,
|
||||
TimezoneType,
|
||||
|
@@ -2,59 +2,36 @@ from functools import wraps
|
||||
|
||||
from sqlalchemy.orm.collections import InstrumentedList as _InstrumentedList
|
||||
|
||||
from .arrow import ArrowType
|
||||
from .choice import Choice, ChoiceType
|
||||
from .color import ColorType
|
||||
from .country import CountryType
|
||||
from .currency import CurrencyType
|
||||
from .email import EmailType
|
||||
from .encrypted import EncryptedType
|
||||
from .ip_address import IPAddressType
|
||||
from .json import JSONType
|
||||
from .locale import LocaleType
|
||||
from .password import Password, PasswordType
|
||||
from .phone_number import PhoneNumber, PhoneNumberType
|
||||
from .range import (
|
||||
from .arrow import ArrowType # noqa
|
||||
from .choice import Choice, ChoiceType # noqa
|
||||
from .color import ColorType # noqa
|
||||
from .country import CountryType # noqa
|
||||
from .currency import CurrencyType # noqa
|
||||
from .email import EmailType # noqa
|
||||
from .encrypted import EncryptedType # noqa
|
||||
from .ip_address import IPAddressType # noqa
|
||||
from .json import JSONType # noqa
|
||||
from .locale import LocaleType # noqa
|
||||
from .password import Password, PasswordType # noqa
|
||||
from .pg_composite import ( # noqa
|
||||
CompositeArray,
|
||||
CompositeType,
|
||||
register_composites,
|
||||
remove_composite_listeners
|
||||
)
|
||||
from .phone_number import PhoneNumber, PhoneNumberType # noqa
|
||||
from .range import ( # noqa
|
||||
DateRangeType,
|
||||
DateTimeRangeType,
|
||||
IntRangeType,
|
||||
NumericRangeType
|
||||
)
|
||||
from .scalar_list import ScalarListException, ScalarListType
|
||||
from .timezone import TimezoneType
|
||||
from .ts_vector import TSVectorType
|
||||
from .url import URLType
|
||||
from .uuid import UUIDType
|
||||
from .weekdays import WeekDaysType
|
||||
|
||||
__all__ = (
|
||||
ArrowType,
|
||||
Choice,
|
||||
ChoiceType,
|
||||
ColorType,
|
||||
CountryType,
|
||||
CurrencyType,
|
||||
DateRangeType,
|
||||
DateTimeRangeType,
|
||||
EmailType,
|
||||
EncryptedType,
|
||||
IntRangeType,
|
||||
IPAddressType,
|
||||
JSONType,
|
||||
LocaleType,
|
||||
NumericRangeType,
|
||||
Password,
|
||||
PasswordType,
|
||||
PhoneNumber,
|
||||
PhoneNumberType,
|
||||
ScalarListException,
|
||||
ScalarListType,
|
||||
TimezoneType,
|
||||
TSVectorType,
|
||||
URLType,
|
||||
UUIDType,
|
||||
WeekDaysType,
|
||||
)
|
||||
from .scalar_list import ScalarListException, ScalarListType # noqa
|
||||
from .timezone import TimezoneType # noqa
|
||||
from .ts_vector import TSVectorType # noqa
|
||||
from .url import URLType # noqa
|
||||
from .uuid import UUIDType # noqa
|
||||
from .weekdays import WeekDaysType # noqa
|
||||
|
||||
|
||||
class InstrumentedList(_InstrumentedList):
|
||||
|
296
sqlalchemy_utils/types/pg_composite.py
Normal file
296
sqlalchemy_utils/types/pg_composite.py
Normal file
@@ -0,0 +1,296 @@
|
||||
"""
|
||||
CompositeType provides means to interact with
|
||||
`PostgreSQL composite types`_. Currently this type features:
|
||||
|
||||
* Easy attribute access to composite type fields
|
||||
* Supports SQLAlchemy TypeDecorator types
|
||||
* Ability to include composite types as part of PostgreSQL arrays
|
||||
* Type creation and dropping
|
||||
|
||||
Installation
|
||||
^^^^^^^^^^^^
|
||||
|
||||
CompositeType automatically attaches `before_create` and `after_drop` DDL
|
||||
listeners. These listeners create and drop the composite type in the
|
||||
database. This means it works out of the box in your test environment where
|
||||
you create the tables on each test run.
|
||||
|
||||
When you already have your database set up you should call
|
||||
:func:`register_composites` after you've set up all models.
|
||||
|
||||
::
|
||||
|
||||
register_composites(conn)
|
||||
|
||||
|
||||
|
||||
Usage
|
||||
^^^^^
|
||||
|
||||
::
|
||||
|
||||
from collections import OrderedDict
|
||||
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy_utils import Composite, CurrencyType
|
||||
|
||||
|
||||
class Account(Base):
|
||||
__tablename__ = 'account'
|
||||
id = sa.Column(sa.Integer, primary_key=True)
|
||||
balance = sa.Column(
|
||||
CompositeType(
|
||||
'money_type',
|
||||
[
|
||||
sa.Column('currency', CurrencyType),
|
||||
sa.Column('amount', sa.Integer)
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
Accessing fields
|
||||
^^^^^^^^^^^^^^^^
|
||||
|
||||
CompositeType provides attribute access to underlying fields. In the following
|
||||
example we find all accounts with balance amount more than 5000.
|
||||
|
||||
|
||||
::
|
||||
|
||||
session.query(Account).filter(Account.balance.amount > 5000)
|
||||
|
||||
|
||||
Arrays of composites
|
||||
^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
::
|
||||
|
||||
from sqlalchemy_utils import CompositeArray
|
||||
|
||||
|
||||
class Account(Base):
|
||||
__tablename__ = 'account'
|
||||
id = sa.Column(sa.Integer, primary_key=True)
|
||||
balances = sa.Column(
|
||||
CompositeArray(
|
||||
CompositeType(
|
||||
'money_type',
|
||||
[
|
||||
sa.Column('currency', CurrencyType),
|
||||
sa.Column('amount', sa.Integer)
|
||||
]
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
.. _PostgreSQL composite types:
|
||||
http://www.postgresql.org/docs/devel/static/rowtypes.html
|
||||
|
||||
|
||||
Related links:
|
||||
|
||||
http://schinckel.net/2014/09/24/using-postgres-composite-types-in-django/
|
||||
"""
|
||||
import psycopg2
|
||||
import sqlalchemy as sa
|
||||
from psycopg2.extensions import adapt, AsIs, register_adapter
|
||||
from sqlalchemy.dialects.postgresql import ARRAY
|
||||
from sqlalchemy.ext.compiler import compiles
|
||||
from sqlalchemy.schema import _CreateDropBase
|
||||
from sqlalchemy.sql.expression import FunctionElement
|
||||
from sqlalchemy.types import (
|
||||
SchemaType,
|
||||
to_instance,
|
||||
TypeDecorator,
|
||||
UserDefinedType
|
||||
)
|
||||
|
||||
|
||||
class CompositeElement(FunctionElement):
|
||||
"""
|
||||
Instances of this class wrap a Postgres composite type.
|
||||
"""
|
||||
def __init__(self, base, field, type_):
|
||||
self.name = field
|
||||
self.type = to_instance(type_)
|
||||
|
||||
super(CompositeElement, self).__init__(base)
|
||||
|
||||
|
||||
@compiles(CompositeElement)
|
||||
def _compile_pgelem(expr, compiler, **kw):
|
||||
return '(%s).%s' % (compiler.process(expr.clauses, **kw), expr.name)
|
||||
|
||||
|
||||
class CompositeArray(ARRAY):
|
||||
def _proc_array(self, arr, itemproc, dim, collection):
|
||||
if dim is None:
|
||||
if issubclass(self.item_type.python_type, (list, tuple)):
|
||||
return arr
|
||||
ARRAY._proc_array(self, arr, itemproc, dim, collection)
|
||||
|
||||
|
||||
# TODO: Make the registration work on connection level instead of global level
|
||||
registered_composites = {}
|
||||
|
||||
|
||||
class CompositeType(UserDefinedType, SchemaType):
|
||||
"""
|
||||
Represents a PostgreSQL composite type.
|
||||
|
||||
:param name:
|
||||
Name of the composite type.
|
||||
:param columns:
|
||||
List of columns that this composite type consists of
|
||||
"""
|
||||
python_type = tuple
|
||||
|
||||
class comparator_factory(UserDefinedType.Comparator):
|
||||
def __getattr__(self, key):
|
||||
try:
|
||||
type_ = self.type.typemap[key]
|
||||
except KeyError:
|
||||
raise KeyError(
|
||||
"Type '%s' doesn't have an attribute: '%s'" % (
|
||||
self.name, key
|
||||
)
|
||||
)
|
||||
|
||||
return CompositeElement(self.expr, key, type_)
|
||||
|
||||
def __init__(self, name, columns):
|
||||
SchemaType.__init__(self)
|
||||
self.name = name
|
||||
self.columns = columns
|
||||
if name in registered_composites:
|
||||
self.type_cls = registered_composites[name].type_cls
|
||||
else:
|
||||
registered_composites[name] = self
|
||||
attach_composite_listeners()
|
||||
|
||||
def get_col_spec(self):
|
||||
return self.name
|
||||
|
||||
def result_processor(self, dialect, coltype):
|
||||
def process(value):
|
||||
cls = value.__class__
|
||||
kwargs = {}
|
||||
for column in self.columns:
|
||||
if isinstance(column.type, TypeDecorator):
|
||||
kwargs[column.name] = column.type.process_result_value(
|
||||
getattr(value, column.name), dialect
|
||||
)
|
||||
else:
|
||||
kwargs[column.name] = getattr(value, column.name)
|
||||
return cls(**kwargs)
|
||||
return process
|
||||
|
||||
def create(self, bind=None, checkfirst=None):
|
||||
if (
|
||||
not checkfirst or
|
||||
not bind.dialect.has_type(bind, self.name, schema=self.schema)
|
||||
):
|
||||
bind.execute(CreateCompositeType(self))
|
||||
|
||||
def drop(self, bind=None, checkfirst=True):
|
||||
if (
|
||||
checkfirst and
|
||||
bind.dialect.has_type(bind, self.name, schema=self.schema)
|
||||
):
|
||||
bind.execute(DropCompositeType(self))
|
||||
|
||||
|
||||
def register_psycopg2_composite(dbapi_connection, composite):
|
||||
composite.type_cls = psycopg2.extras.register_composite(
|
||||
composite.name,
|
||||
dbapi_connection,
|
||||
globally=True
|
||||
).type
|
||||
|
||||
def adapt_composite(value):
|
||||
values = [
|
||||
adapt(getattr(value, column.name)).getquoted().decode('utf-8')
|
||||
for column in
|
||||
composite.columns
|
||||
]
|
||||
return AsIs("(%s)::%s" % (', '.join(values), composite.name))
|
||||
|
||||
register_adapter(composite.type_cls, adapt_composite)
|
||||
|
||||
|
||||
def before_create(target, connection, **kw):
|
||||
for name, composite in registered_composites.items():
|
||||
composite.create(connection, checkfirst=True)
|
||||
register_psycopg2_composite(
|
||||
connection.connection.connection,
|
||||
composite
|
||||
)
|
||||
|
||||
|
||||
def after_drop(target, connection, **kw):
|
||||
for name, composite in registered_composites.items():
|
||||
composite.drop(connection, checkfirst=True)
|
||||
|
||||
|
||||
def register_composites(connection):
|
||||
for name, composite in registered_composites.items():
|
||||
register_psycopg2_composite(
|
||||
connection.connection.connection,
|
||||
composite
|
||||
)
|
||||
|
||||
|
||||
def attach_composite_listeners():
|
||||
listeners = [
|
||||
(sa.MetaData, 'before_create', before_create),
|
||||
(sa.MetaData, 'after_drop', after_drop),
|
||||
]
|
||||
for listener in listeners:
|
||||
if not sa.event.contains(*listener):
|
||||
sa.event.listen(*listener)
|
||||
|
||||
|
||||
def remove_composite_listeners():
|
||||
listeners = [
|
||||
(sa.MetaData, 'before_create', before_create),
|
||||
(sa.MetaData, 'after_drop', after_drop),
|
||||
]
|
||||
for listener in listeners:
|
||||
if sa.event.contains(*listener):
|
||||
sa.event.remove(*listener)
|
||||
|
||||
|
||||
class CreateCompositeType(_CreateDropBase):
|
||||
pass
|
||||
|
||||
|
||||
@compiles(CreateCompositeType)
|
||||
def _visit_create_composite_type(create, compiler, **kw):
|
||||
type_ = create.element
|
||||
fields = ', '.join(
|
||||
'{name} {type}'.format(
|
||||
name=column.name,
|
||||
type=compiler.dialect.type_compiler.process(
|
||||
to_instance(column.type)
|
||||
)
|
||||
)
|
||||
for column in type_.columns
|
||||
)
|
||||
|
||||
return 'CREATE TYPE {name} AS ({fields})'.format(
|
||||
name=compiler.preparer.format_type(type_),
|
||||
fields=fields
|
||||
)
|
||||
|
||||
|
||||
class DropCompositeType(_CreateDropBase):
|
||||
pass
|
||||
|
||||
|
||||
@compiles(DropCompositeType)
|
||||
def _visit_drop_composite_type(drop, compiler, **kw):
|
||||
type_ = drop.element
|
||||
|
||||
return 'DROP TYPE {name}'.format(name=compiler.preparer.format_type(type_))
|
@@ -12,6 +12,7 @@ from sqlalchemy_utils import (
|
||||
i18n,
|
||||
InstrumentedList
|
||||
)
|
||||
from sqlalchemy_utils.types.pg_composite import remove_composite_listeners
|
||||
|
||||
|
||||
@sa.event.listens_for(sa.engine.Engine, 'before_cursor_execute')
|
||||
@@ -60,6 +61,7 @@ class TestCase(object):
|
||||
self.session.close_all()
|
||||
if self.create_tables:
|
||||
self.Base.metadata.drop_all(self.connection)
|
||||
remove_composite_listeners()
|
||||
self.connection.close()
|
||||
self.engine.dispose()
|
||||
|
||||
|
191
tests/types/test_composite.py
Normal file
191
tests/types/test_composite.py
Normal file
@@ -0,0 +1,191 @@
|
||||
import sqlalchemy as sa
|
||||
from pytest import mark
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.ext.declarative import declarative_base
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from sqlalchemy_utils import (
|
||||
CompositeArray,
|
||||
CompositeType,
|
||||
Currency,
|
||||
CurrencyType,
|
||||
i18n,
|
||||
register_composites,
|
||||
remove_composite_listeners
|
||||
)
|
||||
from sqlalchemy_utils.types.currency import babel
|
||||
from tests import TestCase
|
||||
|
||||
|
||||
class TestCompositeTypeWithRegularTypes(TestCase):
|
||||
dns = 'postgres://postgres@localhost/sqlalchemy_utils_test'
|
||||
|
||||
def create_models(self):
|
||||
class Account(self.Base):
|
||||
__tablename__ = 'account'
|
||||
id = sa.Column(sa.Integer, primary_key=True)
|
||||
balance = sa.Column(
|
||||
CompositeType(
|
||||
'money_type',
|
||||
[
|
||||
sa.Column('currency', sa.String),
|
||||
sa.Column('amount', sa.Integer)
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
self.Account = Account
|
||||
|
||||
def test_parameter_processing(self):
|
||||
account = self.Account(
|
||||
balance=('USD', 15)
|
||||
)
|
||||
|
||||
self.session.add(account)
|
||||
self.session.commit()
|
||||
|
||||
account = self.session.query(self.Account).first()
|
||||
assert account.balance.currency == 'USD'
|
||||
assert account.balance.amount == 15
|
||||
|
||||
|
||||
@mark.skipif('babel is None')
|
||||
class TestCompositeTypeWithTypeDecorators(TestCase):
|
||||
dns = 'postgres://postgres@localhost/sqlalchemy_utils_test'
|
||||
|
||||
def setup_method(self, method):
|
||||
TestCase.setup_method(self, method)
|
||||
i18n.get_locale = lambda: babel.Locale('en')
|
||||
|
||||
def create_models(self):
|
||||
class Account(self.Base):
|
||||
__tablename__ = 'account'
|
||||
id = sa.Column(sa.Integer, primary_key=True)
|
||||
balance = sa.Column(
|
||||
CompositeType(
|
||||
'money_type',
|
||||
[
|
||||
sa.Column('currency', CurrencyType),
|
||||
sa.Column('amount', sa.Integer)
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
self.Account = Account
|
||||
|
||||
def test_parameter_processing(self):
|
||||
account = self.Account(
|
||||
balance=('USD', 15)
|
||||
)
|
||||
|
||||
self.session.add(account)
|
||||
self.session.commit()
|
||||
|
||||
account = self.session.query(self.Account).first()
|
||||
assert account.balance.currency == Currency('USD')
|
||||
assert account.balance.amount == 15
|
||||
|
||||
|
||||
@mark.skipif('babel is None')
|
||||
class TestCompositeTypeInsideArray(TestCase):
|
||||
dns = 'postgres://postgres@localhost/sqlalchemy_utils_test'
|
||||
|
||||
def setup_method(self, method):
|
||||
self.type = CompositeType(
|
||||
'money_type',
|
||||
[
|
||||
sa.Column('currency', CurrencyType),
|
||||
sa.Column('amount', sa.Integer)
|
||||
]
|
||||
)
|
||||
|
||||
TestCase.setup_method(self, method)
|
||||
i18n.get_locale = lambda: babel.Locale('en')
|
||||
|
||||
def create_models(self):
|
||||
class Account(self.Base):
|
||||
__tablename__ = 'account'
|
||||
id = sa.Column(sa.Integer, primary_key=True)
|
||||
balances = sa.Column(
|
||||
CompositeArray(self.type)
|
||||
)
|
||||
|
||||
self.Account = Account
|
||||
|
||||
def test_parameter_processing(self):
|
||||
account = self.Account(
|
||||
balances=[
|
||||
self.type.type_cls('USD', 15),
|
||||
self.type.type_cls('AUD', 20)
|
||||
]
|
||||
)
|
||||
|
||||
self.session.add(account)
|
||||
self.session.commit()
|
||||
|
||||
account = self.session.query(self.Account).first()
|
||||
assert account.balances[0].currency == Currency('USD')
|
||||
assert account.balances[0].amount == 15
|
||||
assert account.balances[1].currency == Currency('AUD')
|
||||
assert account.balances[1].amount == 20
|
||||
|
||||
|
||||
class TestCompositeTypeWhenTypeAlreadyExistsInDatabase(TestCase):
|
||||
dns = 'postgres://postgres@localhost/sqlalchemy_utils_test'
|
||||
|
||||
def setup_method(self, method):
|
||||
self.engine = create_engine(self.dns)
|
||||
# self.engine.echo = True
|
||||
self.connection = self.engine.connect()
|
||||
self.Base = declarative_base()
|
||||
|
||||
self.create_models()
|
||||
sa.orm.configure_mappers()
|
||||
|
||||
Session = sessionmaker(bind=self.connection)
|
||||
self.session = Session()
|
||||
self.session.execute(
|
||||
"CREATE TYPE money_type AS (currency VARCHAR, amount INTEGER)"
|
||||
)
|
||||
self.session.execute(
|
||||
"""CREATE TABLE account (
|
||||
id SERIAL, balance MONEY_TYPE, PRIMARY KEY(id)
|
||||
)"""
|
||||
)
|
||||
register_composites(self.connection)
|
||||
|
||||
def teardown_method(self, method):
|
||||
self.session.execute('DROP TABLE account')
|
||||
self.session.execute('DROP TYPE money_type')
|
||||
self.session.close_all()
|
||||
self.connection.close()
|
||||
remove_composite_listeners()
|
||||
self.engine.dispose()
|
||||
|
||||
def create_models(self):
|
||||
class Account(self.Base):
|
||||
__tablename__ = 'account'
|
||||
id = sa.Column(sa.Integer, primary_key=True)
|
||||
balance = sa.Column(
|
||||
CompositeType(
|
||||
'money_type',
|
||||
[
|
||||
sa.Column('currency', sa.String),
|
||||
sa.Column('amount', sa.Integer)
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
self.Account = Account
|
||||
|
||||
def test_parameter_processing(self):
|
||||
account = self.Account(
|
||||
balance=('USD', 15),
|
||||
)
|
||||
|
||||
self.session.add(account)
|
||||
self.session.commit()
|
||||
|
||||
account = self.session.query(self.Account).first()
|
||||
assert account.balance.currency == 'USD'
|
||||
assert account.balance.amount == 15
|
Reference in New Issue
Block a user