How do I use XA in Tortoise-orm?

[tufbel]

I wanted to use XA + Mysql, but so far only find Tortoise.transactions. In_transaction (connection_name=None) in docs, What should I do?

[long2ice]

What's XA?

[tufbel]

What's XA?

XA transaction

[tufbel]

I've written a little bit about the implementation of mysql

from typing import Optional

from pymysql.constants import COMMAND
from tortoise import connections
from tortoise.backends.base.client import TransactionContextPooled
from tortoise.backends.mysql.client import MySQLClient, TransactionWrapper, translate_exceptions
from tortoise.exceptions import ParamsError, TransactionManagementError

class XATransactionWrapper(TransactionWrapper):
    def __init__(self, connection: MySQLClient, transaction_id: str) -> None:
        super().__init__(connection)
        self.transaction_id = transaction_id

    @translate_exceptions
    async def start(self) -> None:
        await self._connection._execute_command(COMMAND.COM_QUERY, f"XA START '{self.transaction_id}'")
        await self._connection._read_ok_packet()
        self._finalized = False

    async def commit(self) -> None:
        if self._finalized:
            raise TransactionManagementError("Transaction already finalised")
        await self._connection._execute_command(COMMAND.COM_QUERY, f"XA END '{self.transaction_id}'")
        await self._connection._read_ok_packet()
        await self._connection._execute_command(COMMAND.COM_QUERY, f"XA PREPARE '{self.transaction_id}'")
        await self._connection._read_ok_packet()
        self._finalized = True

    async def rollback(self) -> None:
        if self._finalized:
            raise TransactionManagementError("Transaction already finalised")
        await self._connection._execute_command(COMMAND.COM_QUERY, f"XA END '{self.transaction_id}'")
        await self._connection._read_ok_packet()
        await self._connection._execute_command(COMMAND.COM_QUERY, f"XA PREPARE '{self.transaction_id}'")
        await self._connection._read_ok_packet()
        self._finalized = True

    async def xa_commit(self) -> None:
        if self._finalized:
            raise TransactionManagementError("Transaction already finalised")
        await self._connection._execute_command(COMMAND.COM_QUERY, f"XA COMMIT '{self.transaction_id}'")
        await self._connection._read_ok_packet()
        self._finalized = True

    async def xa_rollback(self) -> None:
        if self._finalized:
            raise TransactionManagementError("Transaction already finalised")
        await self._connection._execute_command(COMMAND.COM_QUERY, f"XA ROLLBACK '{self.transaction_id}'")
        await self._connection._read_ok_packet()
        self._finalized = True

class XAMySQLClient(MySQLClient):
    def xa_in_transaction(self, transaction_id: str) -> "TransactionContext":
        return TransactionContextPooled(XATransactionWrapper(self, transaction_id))

def _get_connection(connection_name: Optional[str]) -> "BaseDBAsyncClient":
    if connection_name:
        connection = connections.get(connection_name)
    elif len(connections.db_config) == 1:
        connection_name = next(iter(connections.db_config.keys()))
        connection = connections.get(connection_name)
    else:
        raise ParamsError(
            "You are running with multiple databases, so you should specify"
            f" connection_name: {list(connections.db_config)}"
        )
    return connection

def xa_in_transaction(transaction_id: str, connection_name: Optional[str] = None) -> "TransactionContext":
    connection = _get_connection(connection_name)
    return XAMySQLClient.xa_in_transaction(connection, transaction_id)

[long2ice]

In fact I never use that before, of course you can custom a Client for that

[tufbel]

In fact I never use that before, of course you can custom a Client for that

The Client has been rewritten to meet my needs, will Tortoise-orm be considered in a future version?

[long2ice]

Looks like it's not a very common requirement, but PR is welcome