Merge branch 'tortoise:develop' into fk-type-validation

CHANGELOG.rst

@@ -19,10 +19,12 @@ Added
 Fixed
 ^^^^^
 - Fix bug related to `Connector.div` in combined expressions. (#1794)
+- Fix recovery in case of database downtime (#1796)
 
 Changed
 ^^^^^^^
 - Parametrizes UPDATE, DELETE, bulk update and create operations (#1785)
+- Parametrizes related field queries (#1797)
 
 0.22.1
 ------

pyproject.toml

@@ -36,7 +36,7 @@ classifiers = [
 
 [tool.poetry.dependencies]
 python = "^3.8"
-pypika-tortoise = "^0.3.1"
+pypika-tortoise = "^0.3.2"
 iso8601 = "^2.1.0"
 aiosqlite = ">=0.16.0, <0.21.0"
 pytz = "*"

tests/test_transactions.py

@@ -1,4 +1,7 @@
+from unittest.mock import Mock
+
 from tests.testmodels import CharPkModel, Event, Team, Tournament
+from tortoise import connections
 from tortoise.contrib import test
 from tortoise.exceptions import OperationalError, TransactionManagementError
 from tortoise.transactions import atomic, in_transaction
@@ -213,3 +216,27 @@ async def test_select_await_across_transaction_success(self):
         self.assertEqual(
             await Tournament.all().values("id", "name"), [{"id": obj.id, "name": "Test1"}]
         )
+
+
+@test.requireCapability(supports_transactions=True)
+class TestIsolatedTransactions(test.IsolatedTestCase):
+    """Running these in isolation because they mess with the global state of the connections."""
+
+    async def test_rollback_raising_exception(self):
+        """Tests that if a rollback raises an exception, the connection context is restored."""
+        conn = connections.get("models")
+        with self.assertRaisesRegex(ValueError, "rollback"):
+            async with conn._in_transaction() as tx_conn:
+                tx_conn.rollback = Mock(side_effect=ValueError("rollback"))
+                raise ValueError("initial exception")
+
+        self.assertEqual(connections.get("models"), conn)
+
+    async def test_commit_raising_exception(self):
+        """Tests that if a commit raises an exception, the connection context is restored."""
+        conn = connections.get("models")
+        with self.assertRaisesRegex(ValueError, "commit"):
+            async with conn._in_transaction() as tx_conn:
+                tx_conn.commit = Mock(side_effect=ValueError("commit"))
+
+        self.assertEqual(connections.get("models"), conn)

tortoise/backends/base/client.py

@@ -246,7 +246,7 @@ class TransactionContext(Generic[T_conn]):
     def __init__(self, connection: Any) -> None:
         self.connection = connection
         self.connection_name = connection.connection_name
-        self.lock = getattr(connection, "_trxlock", None)
+        self.lock = connection._trxlock
 
     async def ensure_connection(self) -> None:
         if not self.connection._connection:
@@ -255,21 +255,23 @@ async def ensure_connection(self) -> None:
 
     async def __aenter__(self) -> T_conn:
         await self.ensure_connection()
-        await self.lock.acquire()  # type:ignore
+        await self.lock.acquire()
         self.token = connections.set(self.connection_name, self.connection)
         await self.connection.start()
         return self.connection
 
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        if not self.connection._finalized:
-            if exc_type:
-                # Can't rollback a transaction that already failed.
-                if exc_type is not TransactionManagementError:
-                    await self.connection.rollback()
-            else:
-                await self.connection.commit()
-        connections.reset(self.token)
-        self.lock.release()  # type:ignore
+        try:
+            if not self.connection._finalized:
+                if exc_type:
+                    # Can't rollback a transaction that already failed.
+                    if exc_type is not TransactionManagementError:
+                        await self.connection.rollback()
+                else:
+                    await self.connection.commit()
+        finally:
+            connections.reset(self.token)
+            self.lock.release()
 
 
 class TransactionContextPooled(TransactionContext):
@@ -287,16 +289,18 @@ async def __aenter__(self) -> T_conn:
         return self.connection
 
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        if not self.connection._finalized:
-            if exc_type:
-                # Can't rollback a transaction that already failed.
-                if exc_type is not TransactionManagementError:
-                    await self.connection.rollback()
-            else:
-                await self.connection.commit()
-        if self.connection._parent._pool:
-            await self.connection._parent._pool.release(self.connection._connection)
-        connections.reset(self.token)
+        try:
+            if not self.connection._finalized:
+                if exc_type:
+                    # Can't rollback a transaction that already failed.
+                    if exc_type is not TransactionManagementError:
+                        await self.connection.rollback()
+                else:
+                    await self.connection.commit()
+        finally:
+            if self.connection._parent._pool:
+                await self.connection._parent._pool.release(self.connection._connection)
+            connections.reset(self.token)
 
 
 class NestedTransactionContext(TransactionContext):
@@ -313,11 +317,11 @@ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
 
 class NestedTransactionPooledContext(TransactionContext):
     async def __aenter__(self) -> T_conn:
-        await self.lock.acquire()  # type:ignore
+        await self.lock.acquire()
         return self.connection
 
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        self.lock.release()  # type:ignore
+        self.lock.release()
         if not self.connection._finalized:
             if exc_type:
                 # Can't rollback a transaction that already failed.

tortoise/backends/base/executor.py

@@ -21,7 +21,6 @@
 
 from pypika import JoinType, Parameter, Table
 from pypika.queries import QueryBuilder
-from pypika.terms import Parameterizer
 
 from tortoise.exceptions import OperationalError
 from tortoise.expressions import Expression, ResolveContext
@@ -192,10 +191,6 @@ async def _process_insert_result(self, instance: "Model", results: Any) -> None:
     def parameter(self, pos: int) -> Parameter:
         return Parameter(idx=pos + 1)
 
-    @classmethod
-    def parameterizer(cls) -> Parameterizer:
-        return Parameterizer()
-
     async def execute_insert(self, instance: "Model") -> None:
         if not instance._custom_generated_pk:
             values = [
@@ -455,7 +450,7 @@ async def _prefetch_m2m_relation(
             if modifier.having_criterion:
                 query = query.having(modifier.having_criterion)
 
-        _, raw_results = await self.db.execute_query(query.get_sql())
+        _, raw_results = await self.db.execute_query(*query.get_parameterized_sql())
         relations: List[Tuple[Any, Any]] = []
         related_object_list: List["Model"] = []
         model_pk, related_pk = self.model._meta.pk, field_object.related_model._meta.pk

tortoise/backends/psycopg/client.py

@@ -7,6 +7,8 @@
 import psycopg.pq
 import psycopg.rows
 import psycopg_pool
+from pypika.dialects.postgresql import PostgreSQLQuery, PostgreSQLQueryBuilder
+from pypika.terms import Parameterizer
 
 import tortoise.backends.base.client as base_client
 import tortoise.backends.base_postgres.client as postgres_client
@@ -28,7 +30,26 @@ async def release(self, connection: psycopg.AsyncConnection):
         await self.putconn(connection)
 
 
+class PsycopgSQLQuery(PostgreSQLQuery):
+    @classmethod
+    def _builder(cls, **kwargs) -> "PostgreSQLQueryBuilder":
+        return PsycopgSQLQueryBuilder(**kwargs)
+
+
+class PsycopgSQLQueryBuilder(PostgreSQLQueryBuilder):
+    """
+    Psycopg opted to use a custom parameter placeholder, so we need to override the default
+    """
+
+    def get_parameterized_sql(self, **kwargs) -> typing.Tuple[str, list]:
+        parameterizer = kwargs.pop(
+            "parameterizer", Parameterizer(placeholder_factory=lambda _: "%s")
+        )
+        return super().get_parameterized_sql(parameterizer=parameterizer, **kwargs)
+
+
 class PsycopgClient(postgres_client.BasePostgresClient):
+    query_class: typing.Type[PsycopgSQLQuery] = PsycopgSQLQuery
     executor_class: typing.Type[executor.PsycopgExecutor] = executor.PsycopgExecutor
     schema_generator: typing.Type[PsycopgSchemaGenerator] = PsycopgSchemaGenerator
     _pool: typing.Optional[AsyncConnectionPool] = None

tortoise/backends/psycopg/executor.py

@@ -2,7 +2,7 @@
 
 from typing import Optional
 
-from pypika import Parameter, Parameterizer
+from pypika import Parameter
 
 from tortoise import Model
 from tortoise.backends.base_postgres.executor import BasePostgresExecutor
@@ -26,7 +26,3 @@ async def _process_insert_result(
 
     def parameter(self, pos: int) -> Parameter:
         return Parameter("%s")
-
-    @classmethod
-    def parameterizer(cls) -> Parameterizer:
-        return Parameterizer(placeholder_factory=lambda _: "%s")

tortoise/backends/sqlite/executor.py

@@ -1,6 +1,6 @@
 import datetime
-from decimal import Decimal
 import sqlite3
+from decimal import Decimal
 from typing import Optional, Type, Union
 
 import pytz

tortoise/expressions.py

@@ -205,7 +205,8 @@ def __init__(self, query: "AwaitableQuery") -> None:
 
     def get_sql(self, **kwargs: Any) -> str:
         self.query._choose_db_if_not_chosen()
-        return self.query._make_query(**kwargs)[0]
+        self.query._make_query()
+        return self.query.query.get_parameterized_sql(**kwargs)[0]
 
     def as_(self, alias: str) -> "Selectable":  # type: ignore
         self.query._choose_db_if_not_chosen()

tortoise/fields/relational.py

@@ -184,7 +184,9 @@ async def add(self, *instances: MODEL, using_db: "Optional[BaseDBAsyncClient]" =
         criterion = forward_field == pks_f[0] if len(pks_f) == 1 else forward_field.isin(pks_f)
         select_query = select_query.where(criterion)
 
-        _, already_existing_relations_raw = await db.execute_query(str(select_query))
+        _, already_existing_relations_raw = await db.execute_query(
+            *select_query.get_parameterized_sql()
+        )
         already_existing_forward_pks = {
             related_pk_formatting_func(r[forward_key], self.instance)
             for r in already_existing_relations_raw
@@ -194,7 +196,7 @@ async def add(self, *instances: MODEL, using_db: "Optional[BaseDBAsyncClient]" =
             query = db.query_class.into(through_table).columns(forward_field, backward_field)
             for pk_f in pks_f_to_insert:
                 query = query.insert(pk_f, pk_b)
-            await db.execute_query(str(query))
+            await db.execute_query(*query.get_parameterized_sql())
 
     async def clear(self, using_db: "Optional[BaseDBAsyncClient]" = None) -> None:
         """
@@ -237,7 +239,7 @@ async def _remove_or_clear(
                     [related_pk_formatting_func(i.pk, i) for i in instances]
                 )
         query = db.query_class.from_(through_table).where(condition).delete()
-        await db.execute_query(str(query))
+        await db.execute_query(*query.get_parameterized_sql())
 
 
 class RelationalField(Field[MODEL]):