Optimize native versioning

Author: kvesteri

sqlalchemy_continuum/dialects/postgresql.py

@@ -29,6 +29,17 @@
 WHERE NOT EXISTS (SELECT 1 FROM upsert);
 """
 
+temporary_transaction_sql = """
+CREATE TEMP TABLE {temporary_transaction_table}
+({transaction_table_columns})
+ON COMMIT DROP;
+"""
+
+insert_temporary_transaction_sql = """
+INSERT INTO {temporary_transaction_table} ({transaction_table_columns})
+VALUES ({transaction_values});
+"""
+
 procedure_sql = """
 CREATE OR REPLACE FUNCTION {procedure_name}() RETURNS TRIGGER AS $$
 DECLARE transaction_id_value INT;
@@ -40,12 +51,8 @@
             INSERT INTO transaction (native_tx_id)
             VALUES (txid_current()) RETURNING id INTO transaction_id_value;
 
-            CREATE TEMP TABLE continuum_temp_transaction
-            (id BIGINT, PRIMARY KEY(id))
-            ON COMMIT DROP;
-
-            INSERT INTO continuum_temp_transaction (id)
-            VALUES (transaction_id_value);
+            {temporary_transaction_sql}
+            {insert_temporary_transaction_sql}
     END;
 
     IF (TG_OP = 'INSERT') THEN
@@ -127,6 +134,13 @@ def transaction_table_name(self):
         else:
             return 'transaction'
 
+    @property
+    def temporary_transaction_table_name(self):
+        if self.table.schema:
+            return '%s.temporary_transaction' % self.table.schema
+        else:
+            return 'temporary_transaction'
+
     @property
     def version_table_name(self):
         version_table_name = self.version_table_name_format % self.table.name
@@ -345,6 +359,23 @@ def __str__(self):
         )
 
 
+class CreateTemporaryTransactionTableSQL(SQLConstruct):
+    def __str__(self):
+        return temporary_transaction_sql.format(
+            temporary_transaction_table=self.temporary_transaction_table_name,
+            transaction_table_columns='id BIGINT, PRIMARY KEY(id)'
+        )
+
+
+class InsertTemporaryTransactionSQL(SQLConstruct):
+    def __str__(self):
+        return insert_temporary_transaction_sql.format(
+            temporary_transaction_table=self.temporary_transaction_table_name,
+            transaction_table_columns='id',
+            transaction_values='transaction_id_value'
+        )
+
+
 class CreateTriggerFunctionSQL(SQLConstruct):
     def __str__(self):
         args = self.copy_args()
@@ -362,6 +393,12 @@ def __str__(self):
             after_insert=after_insert,
             after_update=after_update,
             after_delete=after_delete,
+            temporary_transaction_sql=(
+                CreateTemporaryTransactionTableSQL(**args)
+            ),
+            insert_temporary_transaction_sql=(
+                InsertTemporaryTransactionSQL(**args)
+            ),
             upsert_insert=InsertUpsertSQL(**args),
             upsert_update=UpdateUpsertSQL(**args),
             upsert_delete=DeleteUpsertSQL(**args)

sqlalchemy_continuum/unit_of_work.py

@@ -1,6 +1,7 @@
 from copy import copy
 
 import sqlalchemy as sa
+from sqlalchemy.orm.attributes import set_committed_value
 from sqlalchemy_utils import get_primary_keys, identity
 from .operation import Operations
 from .utils import (
@@ -11,52 +12,6 @@
     versioned_column_properties
 )
 
-from sqlalchemy.dialects import postgresql
-from sqlalchemy.ext.compiler import compiles
-from sqlalchemy.orm.attributes import set_committed_value
-from sqlalchemy.schema import DDLElement
-from sqlalchemy.sql.expression import CTE, exists
-
-
-class Parenthesis(DDLElement):
-    def __init__(self, element):
-        self.element = element
-
-    def __getattr__(self, attr):
-        return getattr(self.element, attr)
-
-
-@compiles(Parenthesis)
-def visit_alter_column(element, compiler, **kw):
-    return '(%s)' % compiler.process(element.element)
-
-
-def select_or_insert(table, select_values, insert_values, select_criteria):
-    criteria = [
-        getattr(table.c, key) == sa.text(str(value))
-        for key, value in select_criteria.items() if value is not None
-    ]
-
-    select = (
-        sa.select(select_values, from_obj=table)
-        .where(sa.and_(*criteria))
-    )
-    insert = (
-        table.insert()
-        .from_select(
-            insert_values.keys(),
-            sa.select(insert_values.values())
-            .where(~ exists(select))
-        )
-        .returning(*map(sa.text, select_values))
-    )
-    insert_cte = CTE(
-        Parenthesis(insert),
-        name='new_row'
-    )
-    query = sa.select(select_values, from_obj=insert_cte).union(select)
-    return query
-
 
 class UnitOfWork(object):
     def __init__(self, manager):
@@ -159,33 +114,35 @@ def create_transaction(self, session):
         self.current_transaction = Transaction()
 
         if self.manager.options['native_versioning']:
-            criteria = {'native_tx_id': sa.func.txid_current()}
-            args.update(criteria)
-            query = select_or_insert(
-                table,
-                ['*'],
-                args,
-                criteria
-            )
-            query_string = str(query.compile(dialect=postgresql.dialect()))
-
-            values = session.execute(query_string).fetchone()
-            for key, value in values.items():
-                set_committed_value(self.current_transaction, key, value)
-
-            session.execute(
-                '''
-                CREATE TEMP TABLE IF NOT EXISTS continuum_temp_transaction
-                (id BIGINT, PRIMARY KEY(id))
-                ON COMMIT DROP
-                '''
-            )
-            session.execute('''
-                INSERT INTO continuum_temp_transaction (id)
-                SELECT :id WHERE NOT EXISTS
-                (SELECT 1 FROM continuum_temp_transaction WHERE id = :id)''',
-                {'id': self.current_transaction.id}
+            tx_id = (
+                session.execute('SELECT id FROM temporary_transaction')
+                .scalar()
             )
+            if tx_id:
+                set_committed_value(self.current_transaction, 'id', tx_id)
+            else:
+                criteria = {'native_tx_id': sa.func.txid_current()}
+                args.update(criteria)
+
+                query = table.insert().values(**args)
+
+                values = session.execute(query).fetchone()
+                for key, value in values.items():
+                    set_committed_value(self.current_transaction, key, value)
+
+                session.execute(
+                    '''
+                    CREATE TEMP TABLE temporary_transaction
+                    (id BIGINT, PRIMARY KEY(id))
+                    ON COMMIT DROP
+                    '''
+                )
+                session.execute('''
+                    INSERT INTO temporary_transaction (id)
+                    VALUES (:id)
+                    ''',
+                    {'id': self.current_transaction.id}
+                )
             self.merge_transaction(session, self.current_transaction)
         else:
             for key, value in args.items():