Add raw SQL tracking for native versioning

Author: kvesteri

sqlalchemy_continuum/dialects/postgresql.py

@@ -15,10 +15,7 @@
     UPDATE {version_table_name}
     SET {update_values}
     WHERE
-        {transaction_column} = (
-            SELECT MAX(id) FROM {transaction_table_name}
-            WHERE native_tx_id = txid_current()
-        )
+        {transaction_column} = transaction_id_value
         AND
         {primary_key_criteria}
     RETURNING *
@@ -37,7 +34,17 @@
 
 procedure_sql = """
 CREATE OR REPLACE FUNCTION {procedure_name}() RETURNS TRIGGER AS $$
+DECLARE transaction_id_value INT;
 BEGIN
+    transaction_id_value = (
+        SELECT MAX(id) FROM {transaction_table_name}
+        WHERE native_tx_id = txid_current()
+    );
+    IF (transaction_id_value IS NULL) THEN
+        INSERT INTO transaction (native_tx_id)
+        VALUES (txid_current()) RETURNING id INTO transaction_id_value;
+    END IF;
+
     IF (TG_OP = 'INSERT') THEN
         {after_insert}
         {upsert_insert}
@@ -61,10 +68,7 @@
 
 validity_sql = """
 UPDATE {version_table_name}
-SET {end_transaction_column} = (
-    SELECT MAX(id) FROM {transaction_table_name}
-    WHERE native_tx_id = txid_current()
-)
+SET {end_transaction_column} = transaction_id_value
 WHERE
     {transaction_column} = (
         SELECT MIN({transaction_column}) FROM {version_table_name}
@@ -351,6 +355,7 @@ def __str__(self):
             excluded_columns=', '.join(
                 "'%s'" % c for c in self.excluded_columns
             ),
+            transaction_table_name=self.transaction_table_name,
             after_insert=after_insert,
             after_update=after_update,
             after_delete=after_delete,

tests/test_raw_sql.py

@@ -0,0 +1,24 @@
+import pytest
+from tests import TestCase, uses_native_versioning
+
+
+@pytest.mark.skipif('not uses_native_versioning()')
+class TestRawSQL(TestCase):
+    def test_single_statement(self):
+        self.session.execute(
+            "INSERT INTO article (name) VALUES ('some article')"
+        )
+        assert self.session.execute(
+            "SELECT COUNT(1) FROM transaction"
+        ).scalar() == 1
+
+    def test_multiple_statements(self):
+        self.session.execute(
+            "INSERT INTO article (name) VALUES ('some article')"
+        )
+        self.session.execute(
+            "INSERT INTO article (name) VALUES ('some article')"
+        )
+        assert self.session.execute(
+            "SELECT COUNT(1) FROM transaction"
+        ).scalar() == 1