-
Notifications
You must be signed in to change notification settings - Fork 5
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Also use FastAPI dependency for async db transaction.
- Loading branch information
1 parent
9907adb
commit d9bda0b
Showing
7 changed files
with
152 additions
and
15 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,8 @@ | ||
try: | ||
from apphelpers.db.piccolo import * # noqa: F401, F403 | ||
|
||
print("apphelpers: using Piccolo") | ||
except ImportError: | ||
from apphelpers.db.peewee import * # noqa: F401, F403 | ||
|
||
print("apphelpers: using Peewee") |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,55 @@ | ||
from contextlib import asynccontextmanager | ||
from typing import List, Optional, Set, Type, cast | ||
|
||
from piccolo.engine import engine_finder | ||
from piccolo.engine.postgres import PostgresEngine | ||
from piccolo.table import Table, create_db_tables_sync, drop_db_tables_sync | ||
|
||
|
||
@asynccontextmanager | ||
async def connection_pool_lifespan(engine: Optional[PostgresEngine] = None, **kwargs): | ||
if engine := engine or cast(PostgresEngine, engine_finder()): | ||
print("db: starting connection pool") | ||
await engine.start_connection_pool(**kwargs) | ||
yield | ||
print("db: closing connection pool") | ||
await engine.close_connection_pool() | ||
|
||
|
||
@asynccontextmanager | ||
async def dbtransaction_ctx(engine: Optional[PostgresEngine] = None, allow_nested=True): | ||
if engine := engine or cast(PostgresEngine, engine_finder()): | ||
async with engine.transaction(allow_nested=allow_nested): | ||
yield | ||
|
||
|
||
class BaseTable(Table): | ||
@classmethod | ||
def all_column_names(cls) -> Set[str]: | ||
return {col._meta.name for col in cls._meta.columns} | ||
|
||
|
||
def get_sub_tables(basetable: Type[Table]) -> List[type[Table]]: | ||
tables = [] | ||
for subtable in basetable.__subclasses__(): | ||
tables.append(subtable) | ||
tables.extend(get_sub_tables(subtable)) | ||
return tables | ||
|
||
|
||
def setup_db(tables: list[Type[Table]]): | ||
create_db_tables_sync(*tables, if_not_exists=True) | ||
|
||
|
||
def setup_db_from_basetable(basetable: Type[Table]): | ||
tables = get_sub_tables(basetable) | ||
setup_db(tables) | ||
|
||
|
||
def destroy_db(tables: List[Type[Table]]): | ||
drop_db_tables_sync(*tables) | ||
|
||
|
||
def destroy_db_from_basetable(basetable: Type[Table]): | ||
tables = get_sub_tables(basetable) | ||
destroy_db(tables) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -29,3 +29,4 @@ types-pytz | |
types-redis | ||
redis | ||
loguru | ||
piccolo[postgres] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,70 @@ | ||
import asyncio | ||
|
||
import settings | ||
from piccolo import columns as col | ||
from piccolo.engine.postgres import PostgresEngine | ||
|
||
from apphelpers.db.piccolo import ( | ||
BaseTable, | ||
dbtransaction_ctx, | ||
destroy_db_from_basetable, | ||
setup_db_from_basetable, | ||
) | ||
|
||
db = PostgresEngine( | ||
config=dict( | ||
host=settings.DB_HOST, | ||
database=settings.DB_NAME, | ||
user=settings.DB_USER, | ||
password=settings.DB_PASS, | ||
) | ||
) | ||
|
||
|
||
class Book(BaseTable, db=db): | ||
name = col.Text() | ||
|
||
|
||
async def _add_book(name): | ||
await Book.insert(Book(name=name)).run() | ||
|
||
|
||
async def _add_book_loser(name): | ||
await _add_book(name) | ||
loser # will raise # noqa: F821 | ||
|
||
|
||
def setup_module(): | ||
setup_db_from_basetable(BaseTable) | ||
|
||
|
||
def teardown_module(): | ||
destroy_db_from_basetable(BaseTable) | ||
|
||
|
||
async def add_with_tr(): | ||
async with dbtransaction_ctx(db): | ||
name = "The Pillars of the Earth" | ||
await _add_book(name) | ||
names = [b[Book.name] for b in await Book.select().run()] | ||
assert name in names | ||
|
||
try: | ||
async with dbtransaction_ctx(db): | ||
name = "The Cathedral and the Bazaar" | ||
await _add_book_loser(name) | ||
except NameError: | ||
pass | ||
|
||
names = [b[Book.name] for b in await Book.select().run()] | ||
assert name not in names | ||
|
||
async with dbtransaction_ctx(db): | ||
name = "The Ego Trick" | ||
await _add_book(name) | ||
names = [b[Book.name] for b in await Book.select().run()] | ||
assert name in names | ||
|
||
|
||
def test_add_with_tr(): | ||
asyncio.run(add_with_tr()) |