Clean data in db_postgres. :param conn: asyncpg connection object :param table_name_order: list of table name prefixes in the order they should be dropped
(conn, table_name_order: List[str])
| 5 | |
| 6 | |
| 7 | async def clean_data(conn, table_name_order: List[str]): |
| 8 | """ |
| 9 | Clean data in db_postgres. |
| 10 | |
| 11 | :param conn: asyncpg connection object |
| 12 | :param table_name_order: list of table name prefixes in the order they should be dropped |
| 13 | """ |
| 14 | |
| 15 | try: |
| 16 | # Clearing and reinitializing db_postgres |
| 17 | # Implement detailed cleanup steps for tables, sequences, functions, etc. |
| 18 | |
| 19 | # Drop all tables |
| 20 | tables = await conn.fetch("SELECT tablename FROM pg_tables WHERE schemaname='public';") |
| 21 | table_names = [table["tablename"] for table in tables] |
| 22 | |
| 23 | for prefix in table_name_order: |
| 24 | current_table_names = [table_name for table_name in table_names if table_name.startswith(prefix)] |
| 25 | table_names = [table_name for table_name in table_names if table_name not in current_table_names] |
| 26 | for table_name in current_table_names: |
| 27 | await conn.execute(f"DROP TABLE IF EXISTS {table_name} CASCADE;") |
| 28 | |
| 29 | for table_name in table_names: |
| 30 | await conn.execute(f"DROP TABLE IF EXISTS {table_name} CASCADE;") |
| 31 | |
| 32 | # Drop all extensions |
| 33 | extensions = await conn.fetch("""SELECT extname FROM pg_extension;""") |
| 34 | for extension in extensions: |
| 35 | ext_name = extension["extname"] |
| 36 | await conn.execute(f'DROP EXTENSION IF EXISTS "{ext_name}" CASCADE;') |
| 37 | |
| 38 | # Drop all sequences |
| 39 | sequences = await conn.fetch( |
| 40 | "SELECT sequence_name FROM information_schema.sequences WHERE sequence_schema='public';" |
| 41 | ) |
| 42 | for sequence in sequences: |
| 43 | seq_name = sequence["sequence_name"] |
| 44 | await conn.execute(f"DROP SEQUENCE IF EXISTS {seq_name} CASCADE;") |
| 45 | |
| 46 | # Drop all functions |
| 47 | functions = await conn.fetch( |
| 48 | "SELECT routine_name FROM information_schema.routines " |
| 49 | "WHERE specific_schema='public' AND type_udt_name!='trigger';" |
| 50 | ) |
| 51 | for function in functions: |
| 52 | func_name = function["routine_name"] |
| 53 | await conn.execute(f"DROP FUNCTION IF EXISTS {func_name} CASCADE;") |
| 54 | |
| 55 | await conn.execute("CREATE LANGUAGE plpgsql;") |
| 56 | |
| 57 | logger.info("Database cleaned and reinitialized successfully.") |
| 58 | except Exception as e: |
| 59 | logger.error(f"Error during database cleanup: {e}") |