Skip to content

Commit

Permalink
put back last state from main of db engine
Browse files Browse the repository at this point in the history
  • Loading branch information
aakrem committed Jul 10, 2024
1 parent e60409f commit 9fe8df0
Showing 1 changed file with 11 additions and 12 deletions.
23 changes: 11 additions & 12 deletions agenta-backend/agenta_backend/models/db_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -114,20 +114,10 @@ async def initialize_async_postgres(self):
raise ValueError("Postgres URI cannot be None.")

async with self.engine.begin() as conn:
# Drop all existing tables (if needed)
# await conn.run_sync(Base.metadata.drop_all)
# Create tables
# Drop and create tables if needed
for model in models:
await conn.run_sync(model.metadata.create_all)
logger.info(f"Using {self.mode} database...")

async def remove_db(self) -> None:
"""
Remove the database based on the mode.
"""
async with self.engine.begin() as conn:
for model in models:
await conn.run_sync(model.metadata.drop_all)
logger.info(f"Using PostgreSQL database...")

async def initialize_mongodb(self):
"""
Expand Down Expand Up @@ -161,6 +151,15 @@ async def init_db(self):
else:
await self.initialize_async_postgres()

async def remove_db(self) -> None:
"""
Remove the database based on the mode.
"""

async with self.engine.begin() as conn:
for model in models:
await conn.run_sync(model.metadata.drop_all)

@asynccontextmanager
async def get_session(self) -> AsyncGenerator[AsyncSession, None]:
session = self.async_session()
Expand Down

0 comments on commit 9fe8df0

Please sign in to comment.