Source code for aiida.storage.sqlite_zip.utils
###########################################################################
# 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 #
###########################################################################
"""Utilities for this backend."""
import json
import tarfile
import zipfile
from pathlib import Path
from typing import Any, Dict, Optional, Union
from archive_path import read_file_in_tar, read_file_in_zip
from sqlalchemy import event
from sqlalchemy.future.engine import Engine, create_engine
from aiida.common.exceptions import AiidaException, CorruptStorage, UnreachableStorage
META_FILENAME = 'metadata.json'
"""The filename containing meta information about the storage instance."""
DB_FILENAME = 'db.sqlite3'
"""The filename of the SQLite database."""
REPO_FOLDER = 'repo'
"""The name of the folder containing the repository files."""
[docs]
def sqlite_enforce_foreign_keys(dbapi_connection, _):
"""Enforce foreign key constraints, when using sqlite backend (off by default).
See: https://www.sqlite.org/pragma.html#pragma_foreign_keys
"""
cursor = dbapi_connection.cursor()
cursor.execute('PRAGMA foreign_keys=ON;')
cursor.close()
[docs]
def sqlite_case_sensitive_like(dbapi_connection, _):
"""Enforce case sensitive like operations (off by default).
See: https://www.sqlite.org/pragma.html#pragma_case_sensitive_like
"""
cursor = dbapi_connection.cursor()
cursor.execute('PRAGMA case_sensitive_like=ON;')
cursor.close()
[docs]
def create_sqla_engine(path: Union[str, Path], *, enforce_foreign_keys: bool = True, **kwargs) -> Engine:
"""Create a new engine instance."""
engine = create_engine(f'sqlite:///{path}', json_serializer=json.dumps, json_deserializer=json.loads, **kwargs)
event.listen(engine, 'connect', sqlite_case_sensitive_like)
if enforce_foreign_keys:
event.listen(engine, 'connect', sqlite_enforce_foreign_keys)
return engine
[docs]
def read_version(path: Union[str, Path], *, search_limit: Optional[int] = None) -> str:
"""Read the version of the storage instance from the path.
This is intended to work for all versions of the storage format.
:param path: path to storage instance, either a folder, zip file or tar file.
:param search_limit: the maximum number of records to search for the metadata file in a zip file.
:raises: ``UnreachableStorage`` if a version cannot be read from the file
"""
metadata = extract_metadata(path, search_limit=search_limit)
if 'export_version' in metadata:
return metadata['export_version']
raise CorruptStorage("Metadata does not contain 'export_version' key")
[docs]
class ReadOnlyError(AiidaException):
"""Raised when a write operation is called on a read-only archive."""
[docs]
def __init__(self, msg='sqlite_zip storage is read-only'):
super().__init__(msg)