Skip to content

Commit

Permalink
S01E03
Browse files Browse the repository at this point in the history
  • Loading branch information
ansipunk committed Mar 2, 2024
1 parent 6d92566 commit 9d78221
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 5 deletions.
31 changes: 27 additions & 4 deletions databases/backends/psycopg.py
Original file line number Diff line number Diff line change
Expand Up @@ -156,7 +156,7 @@ async def iterate(
yield Record(row, result_columns, self._dialect, column_maps)

def transaction(self) -> "TransactionBackend":
raise NotImplementedError() # pragma: no cover
return PsycopgTransaction(connection=self)

@property
def raw_connection(self) -> typing.Any:
Expand All @@ -166,13 +166,36 @@ def raw_connection(self) -> typing.Any:


class PsycopgTransaction(TransactionBackend):
_connecttion: PsycopgConnection
_transaction: typing.Optional[psycopg.AsyncTransaction]

def __init__(self, connection: PsycopgConnection):
self._connection = connection
self._transaction: typing.Optional[psycopg.AsyncTransaction] = None

async def start(
self, is_root: bool, extra_options: typing.Dict[typing.Any, typing.Any]
) -> None:
raise NotImplementedError() # pragma: no cover
if self._connection._connection is None:
raise RuntimeError("Connection is not acquired")

transaction = psycopg.AsyncTransaction(
self._connection._connection, **extra_options
)
async with transaction._conn.lock:
await transaction._conn.wait(transaction._enter_gen())
self._transaction = transaction

async def commit(self) -> None:
raise NotImplementedError() # pragma: no cover
if self._transaction is None:
raise RuntimeError("Transaction was not started")

async with self._transaction._conn.lock:
await self._transaction._conn.wait(self._transaction._commit_gen())

async def rollback(self) -> None:
raise NotImplementedError() # pragma: no cover
if self._transaction is None:
raise RuntimeError("Transaction was not started")

async with self._transaction._conn.lock:
await self._transaction._conn.wait(self._transaction._rollback_gen(None))
2 changes: 1 addition & 1 deletion databases/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,10 +44,10 @@
class Database:
SUPPORTED_BACKENDS = {
"postgres": "databases.backends.asyncpg:AsyncpgBackend",
"postgresql": "databases.backends.asyncpg:AsyncpgBackend",
"postgresql+aiopg": "databases.backends.aiopg:AiopgBackend",
"postgresql+asyncpg": "databases.backends.asyncpg:AsyncpgBackend",
"postgresql+psycopg": "databases.backends.psycopg:PsycopgBackend",
"postgresql": "databases.backends.psycopg:PsycopgBackend",
"mysql": "databases.backends.mysql:MySQLBackend",
"mysql+aiomysql": "databases.backends.asyncmy:MySQLBackend",
"mysql+asyncmy": "databases.backends.asyncmy:AsyncMyBackend",
Expand Down

0 comments on commit 9d78221

Please sign in to comment.