Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: update load-testing #275

Merged
merged 1 commit into from
Jul 26, 2024
Merged

fix: update load-testing #275

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 17 additions & 25 deletions load-testing/app.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ def new_cursor(self):


logging.basicConfig()
pg_db = peewee_async.PooledPostgresqlDatabase(
database = peewee_async.PooledPostgresqlDatabase(
database='postgres',
user='postgres',
password='postgres',
Expand All @@ -42,36 +42,28 @@ def new_cursor(self):
)


class Manager(peewee_async.Manager):
"""Async models manager."""

database = pg_db


manager = Manager()


def setup_logging():
logger = logging.getLogger("uvicorn.error")
handler = logging.FileHandler(filename="app.log", mode="w")
logger.addHandler(handler)


class MySimplestModel(peewee.Model):
class MySimplestModel(peewee_async.AioModel):
id = peewee.IntegerField(primary_key=True, sequence=True)

class Meta:
database = pg_db
database = database


@asynccontextmanager
async def lifespan(app: FastAPI):
await manager.database.aio_execute_sql('CREATE TABLE IF NOT EXISTS MySimplestModel (id SERIAL PRIMARY KEY);')
await manager.database.aio_execute_sql('TRUNCATE TABLE MySimplestModel;')
await database.aio_execute_sql('CREATE TABLE IF NOT EXISTS MySimplestModel (id SERIAL PRIMARY KEY);')
await database.aio_execute_sql('TRUNCATE TABLE MySimplestModel;')
setup_logging()
yield
# Clean up the ML models and release the resources
await manager.close()
await database.aio_close()

app = FastAPI(lifespan=lifespan)
errors = set()
Expand All @@ -80,28 +72,28 @@ async def lifespan(app: FastAPI):
@app.get("/select")
async def select():
try:
await manager.execute(MySimplestModel.select())
await MySimplestModel.select().aio_execute()
except Exception as e:
errors.add(str(e))
raise
return errors


async def nested_transaction():
async with manager.transaction():
await manager.execute(MySimplestModel.update(id=1))
async with database.aio_atomic():
await MySimplestModel.update(id=1).aio_execute()


async def nested_atomic():
async with manager.database.aio_atomic():
await manager.database.aio_execute(MySimplestModel.update(id=1))
async with database.aio_atomic():
await MySimplestModel.update(id=1).aio_execute()


@app.get("/transaction")
async def transaction():
try:
async with manager.transaction():
await manager.execute(MySimplestModel.update(id=1))
async with database.aio_atomic():
await MySimplestModel.update(id=1).aio_execute()
await nested_transaction()
except Exception as e:
errors.add(str(e))
Expand All @@ -112,8 +104,8 @@ async def transaction():
@app.get("/atomic")
async def atomic():
try:
async with manager.database.aio_atomic():
await manager.database.aio_execute(MySimplestModel.update(id=1))
async with database.aio_atomic():
await MySimplestModel.update(id=1).aio_execute()
await nested_atomic()
except Exception as e:
errors.add(str(e))
Expand All @@ -123,8 +115,8 @@ async def atomic():

@app.get("/recreate_pool")
async def atomic():
await manager.database.aio_close()
await manager.database.aio_connect()
await database.aio_close()
await database.aio_connect()


if __name__ == "__main__":
Expand Down
Loading