From 4e68234033e92d8ec1a591957ba06e87c86f9d9a Mon Sep 17 00:00:00 2001 From: aakrem Date: Tue, 9 Jul 2024 16:34:53 +0200 Subject: [PATCH 1/4] refactor(backend): remove code for deleteign db and tables The code was used when working on migration with main goal is to have a fresh state of tables when we adjust the migration script. --- .../migrations/mongo_to_postgres/db_engine.py | 10 ---------- .../migrations/mongo_to_postgres/migration.py | 1 - .../migrations/mongo_to_postgres/utils.py | 10 ---------- agenta-backend/agenta_backend/models/db_engine.py | 9 --------- 4 files changed, 30 deletions(-) diff --git a/agenta-backend/agenta_backend/migrations/mongo_to_postgres/db_engine.py b/agenta-backend/agenta_backend/migrations/mongo_to_postgres/db_engine.py index 30aa53262..a13dce458 100644 --- a/agenta-backend/agenta_backend/migrations/mongo_to_postgres/db_engine.py +++ b/agenta-backend/agenta_backend/migrations/mongo_to_postgres/db_engine.py @@ -106,21 +106,11 @@ async def init_db(self): Initialize the database based on the mode and create all tables. """ async with self.engine.begin() as conn: - # Drop all existing tables (if needed) - # await conn.run_sync(Base.metadata.drop_all) # Create tables 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) - @asynccontextmanager async def get_session(self) -> AsyncGenerator[AsyncSession, None]: session = self.async_session() diff --git a/agenta-backend/agenta_backend/migrations/mongo_to_postgres/migration.py b/agenta-backend/agenta_backend/migrations/mongo_to_postgres/migration.py index 7154b382d..b3ad69355 100644 --- a/agenta-backend/agenta_backend/migrations/mongo_to_postgres/migration.py +++ b/agenta-backend/agenta_backend/migrations/mongo_to_postgres/migration.py @@ -501,7 +501,6 @@ async def transform_evaluation_scenario(scenario): async def main(): try: - await drop_all_tables() await create_all_tables(tables=tables) await migrate_collection("users", UserDB, transform_user) await migrate_collection("docker_images", ImageDB, transform_image) diff --git a/agenta-backend/agenta_backend/migrations/mongo_to_postgres/utils.py b/agenta-backend/agenta_backend/migrations/mongo_to_postgres/utils.py index 3e5d34545..c93167be6 100644 --- a/agenta-backend/agenta_backend/migrations/mongo_to_postgres/utils.py +++ b/agenta-backend/agenta_backend/migrations/mongo_to_postgres/utils.py @@ -25,16 +25,6 @@ migration_report = {} -async def drop_all_tables(): - """Drop all tables in the database.""" - async with db_engine.engine.begin() as conn: - await conn.run_sync(Base.metadata.reflect) - # Drop all tables with CASCADE option - for table in reversed(Base.metadata.sorted_tables): - await conn.execute(text(f"DROP TABLE IF EXISTS {table.name} CASCADE")) - print("\n====================== All tables are dropped.\n") - - async def create_all_tables(tables): """Create all tables in the database.""" async with db_engine.engine.begin() as conn: diff --git a/agenta-backend/agenta_backend/models/db_engine.py b/agenta-backend/agenta_backend/models/db_engine.py index f9379a969..c8c43da9c 100644 --- a/agenta-backend/agenta_backend/models/db_engine.py +++ b/agenta-backend/agenta_backend/models/db_engine.py @@ -151,15 +151,6 @@ 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() From 47f9cb416e0120ce026184bf916eaf6fdf8380fa Mon Sep 17 00:00:00 2001 From: aakrem Date: Tue, 9 Jul 2024 17:40:54 +0200 Subject: [PATCH 2/4] another missing usage of drop tables --- .../agenta_backend/migrations/mongo_to_postgres/migration.py | 1 - 1 file changed, 1 deletion(-) diff --git a/agenta-backend/agenta_backend/migrations/mongo_to_postgres/migration.py b/agenta-backend/agenta_backend/migrations/mongo_to_postgres/migration.py index b3ad69355..76956c57f 100644 --- a/agenta-backend/agenta_backend/migrations/mongo_to_postgres/migration.py +++ b/agenta-backend/agenta_backend/migrations/mongo_to_postgres/migration.py @@ -34,7 +34,6 @@ ) from agenta_backend.migrations.mongo_to_postgres.utils import ( - drop_all_tables, create_all_tables, print_migration_report, store_mapping, From e60409fb3258bf417e4545e26e3128e939f16b26 Mon Sep 17 00:00:00 2001 From: aakrem Date: Wed, 10 Jul 2024 09:23:42 +0200 Subject: [PATCH 3/4] put back code to be adjusted in the alembic branch --- agenta-backend/agenta_backend/models/db_engine.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/agenta-backend/agenta_backend/models/db_engine.py b/agenta-backend/agenta_backend/models/db_engine.py index c8c43da9c..07436f6a7 100644 --- a/agenta-backend/agenta_backend/models/db_engine.py +++ b/agenta-backend/agenta_backend/models/db_engine.py @@ -114,10 +114,20 @@ async def initialize_async_postgres(self): raise ValueError("Postgres URI cannot be None.") async with self.engine.begin() as conn: - # Drop and create tables if needed + # Drop all existing tables (if needed) + # await conn.run_sync(Base.metadata.drop_all) + # Create tables for model in models: await conn.run_sync(model.metadata.create_all) - logger.info(f"Using PostgreSQL database...") + 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) async def initialize_mongodb(self): """ From 9fe8df0799e20b417f383f08b5535fac1be193f9 Mon Sep 17 00:00:00 2001 From: aakrem Date: Wed, 10 Jul 2024 09:27:28 +0200 Subject: [PATCH 4/4] put back last state from main of db engine --- .../agenta_backend/models/db_engine.py | 23 +++++++++---------- 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/agenta-backend/agenta_backend/models/db_engine.py b/agenta-backend/agenta_backend/models/db_engine.py index 07436f6a7..f9379a969 100644 --- a/agenta-backend/agenta_backend/models/db_engine.py +++ b/agenta-backend/agenta_backend/models/db_engine.py @@ -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): """ @@ -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()