Merge pull request #14 from CloverHealth/joined-table-inheritance

Joined table inheritance

Author: multigl

requirements.txt

@@ -1,2 +1,2 @@
-sqlalchemy==1.0.15
+sqlalchemy==1.1.2
 psycopg2==2.6.2

temporal_sqlalchemy/clock.py

@@ -124,7 +124,7 @@ def make_temporal(cls: nine.Type[Clocked]):
         clock_table_name = _truncate_identifier("%s_clock" % entity_table_name)
 
         history_tables = {
-            p: build_history_class(p, schema)
+            p: build_history_class(cls, p, schema)
             for p in local_props | relationship_props
         }
 
@@ -221,12 +221,12 @@ def build_clock_class(
 
 
 def build_history_class(
+        cls: declarative.DeclarativeMeta,
         prop: T_PROPS,
         schema: str = None) -> nine.Type[TemporalProperty]:
     """build a sql alchemy table for given prop"""
-    cls = prop.parent.class_
     class_name = "%s%s_%s" % (cls.__name__, 'History', prop.key)
-    table = build_history_table(prop, schema)
+    table = build_history_table(cls, prop, schema)
     base_classes = (
         TemporalProperty,
         declarative.declarative_base(metadata=cls.metadata),
@@ -252,7 +252,10 @@ def build_history_class(
     return model
 
 
-def build_history_table(prop: T_PROPS, schema: str = None) -> sa.Table:
+def build_history_table(
+        cls: declarative.DeclarativeMeta,
+        prop: T_PROPS,
+        schema: str = None) -> sa.Table:
     """build a sql alchemy table for given prop"""
 
     if isinstance(prop, orm.RelationshipProperty):
@@ -267,7 +270,7 @@ def build_history_table(prop: T_PROPS, schema: str = None) -> sa.Table:
         property_key = prop.key
         columns = (_copy_column(col) for col in prop.columns)
 
-    local_table = prop.parent.local_table
+    local_table = cls.__table__
     table_name = _truncate_identifier(
         '%s_%s_%s' % (local_table.name, 'history', property_key))
     index_name = _truncate_identifier('%s_effective_idx' % table_name)
@@ -289,7 +292,7 @@ def build_history_table(prop: T_PROPS, schema: str = None) -> sa.Table:
     ]
 
     # TODO: make this support different shape pks
-    foreign_key = getattr(prop.parent.class_, 'id')
+    foreign_key = getattr(cls, 'id')
     return sa.Table(
         table_name,
         prop.parent.class_.metadata,

temporal_sqlalchemy/version.py

@@ -1,2 +1,2 @@
 """Version information."""
-__version__ = '0.1.0'
+__version__ = '0.2.0'

tests/models.py

@@ -165,3 +165,42 @@ class HugeIndices(temporal_sqlalchemy.Clocked, EdgeCaseBase):
 
     id = auto_uuid()
     really_really_really_really_really_long_column = sa.Column(sa.Integer)
+
+
+class JoinedEnumBase(Base):
+    __tablename__ = 'joined_enum_base'
+
+    id = auto_uuid()
+    kind = sa.Column(
+        sap.ENUM('default', 'enum_a', 'enum_b', name='joined_enum_kind'))
+    is_deleted = sa.Column(sa.Boolean, default=False)
+
+    __mapper_args__ = {
+        'polymorphic_on': kind,
+        'polymorphic_identity': 'default'
+    }
+
+
+@temporal_sqlalchemy.add_clock(
+    'val',
+    'is_deleted',
+    temporal_schema=TEMPORAL_SCHEMA)
+class JoinedEnumA(temporal_sqlalchemy.Clocked, JoinedEnumBase):
+    __tablename__ = 'joined_enum_a'
+
+    id = sa.Column(sa.ForeignKey(JoinedEnumBase.id), primary_key=True)
+    val = sa.Column(sap.ENUM('foo', 'foobar', name='joined_enum_a_val'))
+
+    __mapper_args__ = {'polymorphic_identity': 'enum_a'}
+
+
+@temporal_sqlalchemy.add_clock(
+    'val',
+    'is_deleted', temporal_schema=TEMPORAL_SCHEMA)
+class JoinedEnumB(temporal_sqlalchemy.Clocked, JoinedEnumBase):
+    __tablename__ = 'joined_enum_b'
+
+    id = sa.Column(sa.ForeignKey(JoinedEnumBase.id), primary_key=True)
+    val = sa.Column(sap.ENUM('bar', 'barfoo', name='joined_enum_b_val'))
+
+    __mapper_args__ = {'polymorphic_identity': 'enum_b'}

tests/test_builders.py

@@ -11,10 +11,13 @@ def test_build_history_table():
     rel_id_prop = sa.inspect(models.RelationalTemporalModel.rel_id).property
     rel_prop = sa.inspect(models.RelationalTemporalModel.rel).property
 
-    history_table = build_history_table(rel_id_prop, models.TEMPORAL_SCHEMA)
+    history_table = build_history_table(
+        models.RelationalTemporalModel, rel_id_prop, models.TEMPORAL_SCHEMA)
 
     assert history_table == build_history_table(
-        rel_prop, models.TEMPORAL_SCHEMA)
+        models.RelationalTemporalModel,
+        rel_prop,
+        models.TEMPORAL_SCHEMA)
     assert history_table.name == 'relational_temporal_history_rel_id'
     assert history_table.schema == models.TEMPORAL_SCHEMA
     assert history_table.c.keys() == [
@@ -26,8 +29,8 @@ def test_build_history_class():
     rel_id_prop = sa.inspect(models.SimpleTable.rel_id).property
     rel_prop = sa.inspect(models.SimpleTable.rel).property
 
-    rel_id_prop_class = build_history_class(rel_id_prop)
-    rel_prop_class = build_history_class(rel_prop)
+    rel_id_prop_class = build_history_class(models.SimpleTable, rel_id_prop)
+    rel_prop_class = build_history_class(models.SimpleTable, rel_prop)
 
     assert rel_id_prop_class.__table__ == rel_prop_class.__table__
     assert rel_id_prop_class.__name__ == 'SimpleTableHistory_rel_id'

tests/test_joined_table_inheritance.py

@@ -0,0 +1,71 @@
+import pytest
+import sqlalchemy as sa
+
+import temporal_sqlalchemy as temporal
+
+from . import models
+
+
+@pytest.fixture(autouse=True)
+def setup(session):
+    """Generate schema and tables from a list of model bases."""
+    for schema in models.basic_metadata._schemas:
+        session.execute('CREATE SCHEMA IF NOT EXISTS %s' % schema)
+
+    models.basic_metadata.create_all(session.bind)
+
+
+def test_joined_enums_create(session):
+    session.add_all([
+        models.JoinedEnumA(val='foo'),
+        models.JoinedEnumA(val='foobar'),
+        models.JoinedEnumB(val='bar'),
+        models.JoinedEnumB(val='barfoo'),
+        models.JoinedEnumB(val='barfoo'),
+    ])
+
+    session.commit()
+
+    assert session.query(models.JoinedEnumBase).count() == 5
+
+    enum_a_val_history = temporal.get_history_model(models.JoinedEnumA.val)
+    assert session.query(enum_a_val_history).count() == 2
+
+    enum_b_val_history = temporal.get_history_model(models.JoinedEnumB.val)
+    assert session.query(enum_b_val_history).count() == 3
+
+
+@pytest.mark.parametrize("model,first_val,second_val", (
+        (models.JoinedEnumA, 'foo', 'foobar'),
+        (models.JoinedEnumA, 'foobar', 'foo'),
+        (models.JoinedEnumB, 'bar', 'barfoo'),
+))
+def test_joined_enums_edit(session, model, first_val, second_val):
+    kind = sa.inspect(model).polymorphic_identity
+
+    session.add(model(val=first_val, is_deleted=False))
+    session.commit()
+
+    entity = session.query(model).first()
+    entity_id = entity.id
+    with entity.clock_tick():
+        entity.is_deleted = True
+        entity.val = second_val
+
+    session.commit()
+    session.expunge_all()  # empty the session
+
+    # query directly from the model
+    entity = session.query(model).get(entity_id)
+    assert entity.vclock == 2
+    assert entity.val == second_val
+    assert entity.kind == kind
+    assert entity.is_deleted is True
+
+    # query via with_polymorphic
+    entity = session.query(models.JoinedEnumBase)\
+        .with_polymorphic(model).first()
+    assert entity.vclock == 2
+    assert entity.val == second_val
+    assert entity.kind == kind
+    assert entity.is_deleted is True

tests/test_temporal_multi_entity_with_activity.py

@@ -1,6 +1,4 @@
 import pytest
-import sqlalchemy.exc as exc
-import sqlalchemy.orm.attributes as orm_attr
 
 import temporal_sqlalchemy as temporal
 
@@ -11,7 +9,7 @@ class TestTemporalMultiEntityWithActivity(shared.DatabaseTest):
     @pytest.fixture(autouse=True)
     def setup(self, session):
         models.activity_metadata.create_all(session.bind)
-    
+
     def test_activity_on_multi_entity_create(self, session):
         activity = models.Activity(description='Create temps')
         session.add(activity)
@@ -42,9 +40,11 @@ def test_activity_on_multi_entity_create(self, session):
         assert t2.vclock == 1
         assert t2.clock.count() == 1
 
-        clock1_query = session.query(models.FirstTemporalWithActivity.temporal_options.clock_table)
+        clock1_query = session.query(
+            models.FirstTemporalWithActivity.temporal_options.clock_table)
         assert clock1_query.count() == 1
-        clock2_query = session.query(models.SecondTemporalWithActivity.temporal_options.clock_table)
+        clock2_query = session.query(
+            models.SecondTemporalWithActivity.temporal_options.clock_table)
         assert clock2_query.count() == 1
 
         clock1_result = clock1_query.first()
@@ -57,8 +57,10 @@ def test_activity_on_multi_entity_edit(self, session):
         create_activity = models.Activity(description='Create temp')
         session.add(create_activity)
 
-        t1 = models.FirstTemporalWithActivity(column=1234, activity=create_activity)
-        t2 = models.SecondTemporalWithActivity(column=4567, activity=create_activity)
+        t1 = models.FirstTemporalWithActivity(
+            column=1234, activity=create_activity)
+        t2 = models.SecondTemporalWithActivity(
+            column=4567, activity=create_activity)
         session.add(t1)
         session.add(t2)
         session.commit()
@@ -77,7 +79,9 @@ def test_activity_on_multi_entity_edit(self, session):
         activity_query = session.query(models.Activity)
         assert activity_query.count() == 2
 
-        create_activity_result = activity_query.order_by(models.Activity.date_created).first()
+        create_activity_result = activity_query\
+            .order_by(models.Activity.date_created)\
+            .first()
         assert create_activity_result.description == 'Create temp'
 
         activity_clock2_backref = temporal.get_activity_clock_backref(
@@ -88,7 +92,9 @@ def test_activity_on_multi_entity_edit(self, session):
         assert getattr(create_activity_result, activity_clock2_backref.key)
         assert getattr(create_activity_result, activity_clock1_backref.key)
 
-        edit_activity_result = activity_query.order_by(models.Activity.date_created.desc()).first()
+        edit_activity_result = activity_query\
+            .order_by(models.Activity.date_created.desc())\
+            .first()
         assert edit_activity_result.description == 'Edit temp'
 
         assert getattr(edit_activity, activity_clock2_backref.key)
@@ -102,15 +108,19 @@ def test_activity_on_multi_entity_edit(self, session):
         assert t2.vclock == 2
         assert t2.clock.count() == 2
 
-        clock1_query = session.query(models.FirstTemporalWithActivity.temporal_options.clock_table)\
-            .order_by(models.FirstTemporalWithActivity.temporal_options.clock_table.tick).all()
+        clock_model = models.FirstTemporalWithActivity.temporal_options\
+            .clock_table
+        clock1_query = session.query(clock_model)\
+            .order_by(clock_model.tick).all()
         assert len(clock1_query) == 2
 
         assert clock1_query[0].activity_id == create_activity_result.id
         assert clock1_query[1].activity_id == edit_activity_result.id
 
-        clock2_query = session.query(models.SecondTemporalWithActivity.temporal_options.clock_table)\
-            .order_by(models.SecondTemporalWithActivity.temporal_options.clock_table.tick).all()
+        clock_model = models.SecondTemporalWithActivity.temporal_options\
+            .clock_table
+        clock2_query = session.query(clock_model)\
+            .order_by(clock_model.tick).all()
         assert len(clock2_query) == 2
 
         assert clock2_query[0].activity_id == create_activity_result.id