Source code for aiida.storage.psql_dos.migrations.utils.dblog_update

# -*- coding: utf-8 -*-
###########################################################################
# Copyright (c), The AiiDA team. All rights reserved.                     #
# This file is part of the AiiDA code.                                    #
#                                                                         #
# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
# For further information on the license, see the LICENSE.txt file        #
# For further information please visit http://www.aiida.net               #
###########################################################################
"""Shared function for django_0024 and sqlalchemy ea2f50e7f615"""
import sys
from tempfile import NamedTemporaryFile
from typing import Set

import click
import sqlalchemy as sa

from aiida.cmdline.utils import echo

from .utils import dumps_json


[docs]def get_legacy_workflow_log_number(connection): """ Get the number of the log records that correspond to legacy workflows """ return connection.execute( sa.text( """ SELECT COUNT(*) FROM db_dblog WHERE (db_dblog.objname LIKE 'aiida.workflows.user.%') """ ) ).fetchall()[0][0]
[docs]def get_unknown_entity_log_number(connection): """ Get the number of the log records that correspond to unknown entities """ return connection.execute( sa.text( """ SELECT COUNT(*) FROM db_dblog WHERE (db_dblog.objname NOT LIKE 'node.%') AND (db_dblog.objname NOT LIKE 'aiida.workflows.user.%') """ ) ).fetchall()[0][0]
[docs]def get_logs_with_no_nodes_number(connection): """ Get the number of the log records that correspond to nodes that were deleted """ return connection.execute( sa.text( """ SELECT COUNT(*) FROM db_dblog WHERE (db_dblog.objname LIKE 'node.%') AND NOT EXISTS (SELECT 1 FROM db_dbnode WHERE db_dbnode.id = db_dblog.objpk LIMIT 1) """ ) ).fetchall()[0][0]
[docs]def get_serialized_legacy_workflow_logs(connection): """ Get the serialized log records that correspond to legacy workflows """ query = connection.execute( sa.text( """ SELECT db_dblog.id, db_dblog.time, db_dblog.loggername, db_dblog.levelname, db_dblog.objpk, db_dblog.objname, db_dblog.message, db_dblog.metadata FROM db_dblog WHERE (db_dblog.objname LIKE 'aiida.workflows.user.%') """ ) ) res = [] for row in query: res.append(row._asdict()) return dumps_json(res)
[docs]def get_serialized_unknown_entity_logs(connection): """ Get the serialized log records that correspond to unknown entities """ query = connection.execute( sa.text( """ SELECT db_dblog.id, db_dblog.time, db_dblog.loggername, db_dblog.levelname, db_dblog.objpk, db_dblog.objname, db_dblog.message, db_dblog.metadata FROM db_dblog WHERE (db_dblog.objname NOT LIKE 'node.%') AND (db_dblog.objname NOT LIKE 'aiida.workflows.user.%') """ ) ) res = [] for row in query: res.append(row._asdict()) return dumps_json(res)
[docs]def get_serialized_logs_with_no_nodes(connection): """ Get the serialized log records that correspond to nodes that were deleted """ query = connection.execute( sa.text( """ SELECT db_dblog.id, db_dblog.time, db_dblog.loggername, db_dblog.levelname, db_dblog.objpk, db_dblog.objname, db_dblog.message, db_dblog.metadata FROM db_dblog WHERE (db_dblog.objname LIKE 'node.%') AND NOT EXISTS (SELECT 1 FROM db_dbnode WHERE db_dbnode.id = db_dblog.objpk LIMIT 1) """ ) ) res = [] for row in query: res.append(row._asdict()) return dumps_json(res)
[docs]def export_and_clean_workflow_logs(connection, profile): """Export the logs records that correspond to legacy workflows and to unknown entities (place them to files and remove them from the DbLog table). """ lwf_no_number = get_legacy_workflow_log_number(connection) other_number = get_unknown_entity_log_number(connection) log_no_node_number = get_logs_with_no_nodes_number(connection) # If there are no legacy workflow log records or log records of an unknown entity if lwf_no_number == 0 and other_number == 0 and log_no_node_number == 0: return if not profile.is_test_profile: echo.echo_warning( 'We found {} log records that correspond to legacy workflows and {} log records to correspond ' 'to an unknown entity.'.format(lwf_no_number, other_number) ) echo.echo_warning( 'These records will be removed from the database and exported to JSON files (to the current directory).' ) proceed = click.confirm('Would you like to proceed?', default=True) if not proceed: sys.exit(1) delete_on_close = profile.is_test_profile # Exporting the legacy workflow log records if lwf_no_number != 0: # Get the records and write them to file with NamedTemporaryFile( prefix='legagy_wf_logs-', suffix='.log', dir='.', delete=delete_on_close, mode='w+' ) as handle: # Export the log records filename = handle.name handle.write(get_serialized_legacy_workflow_logs(connection)) # If delete_on_close is False, we are running for the user and add additional message of file location if not delete_on_close: echo.echo(f'Exported legacy workflow logs to {filename}') # Now delete the records connection.execute( sa.text( """ DELETE FROM db_dblog WHERE (db_dblog.objname LIKE 'aiida.workflows.user.%') """ ) ) # Exporting unknown log records if other_number != 0: # Get the records and write them to file with NamedTemporaryFile( prefix='unknown_entity_logs-', suffix='.log', dir='.', delete=delete_on_close, mode='w+' ) as handle: # Export the log records filename = handle.name handle.write(get_serialized_unknown_entity_logs(connection)) # If delete_on_close is False, we are running for the user and add additional message of file location if not delete_on_close: echo.echo(f'Exported unexpected entity logs to {filename}') # Now delete the records connection.execute( sa.text( """ DELETE FROM db_dblog WHERE (db_dblog.objname NOT LIKE 'node.%') AND (db_dblog.objname NOT LIKE 'aiida.workflows.user.%') """ ) ) # Exporting log records that don't correspond to nodes if log_no_node_number != 0: # Get the records and write them to file with NamedTemporaryFile( prefix='no_node_entity_logs-', suffix='.log', dir='.', delete=delete_on_close, mode='w+' ) as handle: # Export the log records filename = handle.name handle.write(get_serialized_logs_with_no_nodes(connection)) # If delete_on_close is False, we are running for the user and add additional message of file location if not delete_on_close: echo.echo(f'Exported entity logs that do not correspond to nodes to {filename}') # Now delete the records connection.execute( sa.text( """ DELETE FROM db_dblog WHERE (db_dblog.objname LIKE 'node.%') AND NOT EXISTS (SELECT 1 FROM db_dbnode WHERE db_dbnode.id = db_dblog.objpk LIMIT 1) """ ) )
[docs]def set_new_uuid(connection): """ Set new and distinct UUIDs to all the logs """ from aiida.common.utils import get_new_uuid # Exit if there are no rows - e.g. initial setup id_query = connection.execute(sa.text('SELECT db_dblog.id FROM db_dblog')) if id_query.rowcount == 0: return id_res = id_query.fetchall() ids = [] for (curr_id,) in id_res: ids.append(curr_id) uuids: Set[str] = set() while len(uuids) < len(ids): uuids.add(get_new_uuid()) # Create the key/value pairs key_values = ','.join(f"({curr_id}, '{curr_uuid}')" for curr_id, curr_uuid in zip(ids, uuids)) update_stm = f""" UPDATE db_dblog as t SET uuid = uuid(c.uuid) from (values {key_values}) as c(id, uuid) where c.id = t.id""" connection.execute(sa.text(update_stm))