From 951578741b2e6576e9b2deab2e6a06532e1a77e3 Mon Sep 17 00:00:00 2001 From: Sebastiaan Huber Date: Wed, 16 Nov 2022 13:15:13 +0100 Subject: [PATCH 1/2] `PsqlDosMigrator`: Remove hardcoding of table name in database reset The `PsqlDosMigrator.reset_database` and `PsqlDosBackend._clear` methods both deleted (almost) all database tables. The tables to delete were hard-coded in a list by each method. The `delete_tables` method is added to the `PsqlDosMigrator` class which will delete all tables of the schema which are determined dynamically by reflecting the metadata schema. The `exclude_tables` argument can be used to define table names to exclude. The new method is used by both classes in order to remove the hard-coding of table names. --- aiida/storage/psql_dos/backend.py | 20 ++++++++++--------- aiida/storage/psql_dos/migrator.py | 31 ++++++++++++++++++++++-------- 2 files changed, 34 insertions(+), 17 deletions(-) diff --git a/aiida/storage/psql_dos/backend.py b/aiida/storage/psql_dos/backend.py index 709acd2ba9..bc7cbf4b26 100644 --- a/aiida/storage/psql_dos/backend.py +++ b/aiida/storage/psql_dos/backend.py @@ -16,7 +16,6 @@ from typing import TYPE_CHECKING, Iterator, List, Optional, Sequence, Set, Union from disk_objectstore import Container -from sqlalchemy import table from sqlalchemy.orm import Session, scoped_session, sessionmaker from aiida.common.exceptions import ClosedStorage, ConfigurationError, IntegrityError @@ -170,18 +169,21 @@ def _clear(self) -> None: super()._clear() - session = self.get_session() - with self.migrator_context(self._profile) as migrator: - with self.transaction(): - for table_name in ( - 'db_dbgroup_dbnodes', 'db_dbgroup', 'db_dblink', 'db_dbnode', 'db_dblog', 'db_dbauthinfo', - 'db_dbuser', 'db_dbcomputer' - ): - session.execute(table(table_name).delete()) + # First clear the contents of the database + with self.transaction() as session: + + # Close the session otherwise the ``delete_tables`` call will hang as there will be an open connection + # to the PostgreSQL server and it will block the deletion and the command will hang. + self.get_session().close() + exclude_tables = [migrator.alembic_version_tbl_name, 'db_dbsetting'] + migrator.delete_tables(exclude_tables) + + # Clear out all references to database model instances which are now invalid. session.expunge_all() + # Now reset and reinitialise the repository migrator.reset_repository() migrator.initialise_repository() repository_uuid = migrator.get_repository_uuid() diff --git a/aiida/storage/psql_dos/migrator.py b/aiida/storage/psql_dos/migrator.py index 1c406eb329..d5563e0f94 100644 --- a/aiida/storage/psql_dos/migrator.py +++ b/aiida/storage/psql_dos/migrator.py @@ -27,7 +27,7 @@ from alembic.runtime.migration import MigrationContext, MigrationInfo from alembic.script import ScriptDirectory from disk_objectstore import Container -from sqlalchemy import String, Table, column, desc, insert, inspect, select, table +from sqlalchemy import MetaData, String, Table, column, desc, insert, inspect, select, table from sqlalchemy.exc import OperationalError, ProgrammingError from sqlalchemy.ext.automap import automap_base from sqlalchemy.orm import Session @@ -272,13 +272,7 @@ def reset_database(self) -> None: This will also destroy the settings table and so in order to use it again, it will have to be reinitialised. """ - if inspect(self.connection).has_table(self.alembic_version_tbl_name): - for table_name in ( - 'db_dbgroup_dbnodes', 'db_dbgroup', 'db_dblink', 'db_dbnode', 'db_dblog', 'db_dbauthinfo', 'db_dbuser', - 'db_dbcomputer', 'db_dbsetting' - ): - self.connection.execute(table(table_name).delete()) - self.connection.commit() + self.delete_tables(exclude_tables=[self.alembic_version_tbl_name]) def initialise_repository(self) -> None: """Initialise the repository.""" @@ -312,6 +306,27 @@ def initialise_database(self) -> None: context.stamp(context.script, 'main@head') self.connection.commit() + def delete_tables(self, exclude_tables: list[str] | None = None) -> None: + """Delete the tables of the database. + + :param exclude_tables: Optional list of table names that should not be deleted. + """ + exclude_tables = exclude_tables or [] + + if inspect(self.connection).has_table(self.alembic_version_tbl_name): + + metadata = MetaData() + metadata.reflect(bind=self.connection) + + # The ``sorted_tables`` property returns the tables sorted by their foreign-key dependencies, with those + # that are dependent on others first. Iterate over the list in reverse to ensure that the tables with + # the independent rows first. + for schema_table in reversed(metadata.sorted_tables): + if schema_table.name in exclude_tables: + continue + self.connection.execute(schema_table.delete()) + self.connection.commit() + def migrate(self) -> None: """Migrate the storage for this profile to the head version. From ad1250c3746827b4aa35675f6d055e8da6c1fd9a Mon Sep 17 00:00:00 2001 From: Sebastiaan Huber Date: Thu, 24 Nov 2022 10:03:50 +0100 Subject: [PATCH 2/2] Address review comments --- aiida/storage/psql_dos/backend.py | 2 +- aiida/storage/psql_dos/migrator.py | 11 +++++++---- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/aiida/storage/psql_dos/backend.py b/aiida/storage/psql_dos/backend.py index bc7cbf4b26..57e0686d3a 100644 --- a/aiida/storage/psql_dos/backend.py +++ b/aiida/storage/psql_dos/backend.py @@ -178,7 +178,7 @@ def _clear(self) -> None: # to the PostgreSQL server and it will block the deletion and the command will hang. self.get_session().close() exclude_tables = [migrator.alembic_version_tbl_name, 'db_dbsetting'] - migrator.delete_tables(exclude_tables) + migrator.delete_all_tables(exclude_tables=exclude_tables) # Clear out all references to database model instances which are now invalid. session.expunge_all() diff --git a/aiida/storage/psql_dos/migrator.py b/aiida/storage/psql_dos/migrator.py index d5563e0f94..47156b1d29 100644 --- a/aiida/storage/psql_dos/migrator.py +++ b/aiida/storage/psql_dos/migrator.py @@ -272,7 +272,7 @@ def reset_database(self) -> None: This will also destroy the settings table and so in order to use it again, it will have to be reinitialised. """ - self.delete_tables(exclude_tables=[self.alembic_version_tbl_name]) + self.delete_all_tables(exclude_tables=[self.alembic_version_tbl_name]) def initialise_repository(self) -> None: """Initialise the repository.""" @@ -306,8 +306,11 @@ def initialise_database(self) -> None: context.stamp(context.script, 'main@head') self.connection.commit() - def delete_tables(self, exclude_tables: list[str] | None = None) -> None: - """Delete the tables of the database. + def delete_all_tables(self, *, exclude_tables: list[str] | None = None) -> None: + """Delete all tables of the current database schema. + + The tables are determined dynamically through reflection of the current schema version. Any other tables in the + database that are not part of the schema should remain unaffected. :param exclude_tables: Optional list of table names that should not be deleted. """ @@ -320,7 +323,7 @@ def delete_tables(self, exclude_tables: list[str] | None = None) -> None: # The ``sorted_tables`` property returns the tables sorted by their foreign-key dependencies, with those # that are dependent on others first. Iterate over the list in reverse to ensure that the tables with - # the independent rows first. + # the independent rows are deleted first. for schema_table in reversed(metadata.sorted_tables): if schema_table.name in exclude_tables: continue