| # Copyright (c) 2019-2021 The Regents of the University of California |
| # All Rights Reserved. |
| # |
| # Redistribution and use in source and binary forms, with or without |
| # modification, are permitted provided that the following conditions are |
| # met: redistributions of source code must retain the above copyright |
| # notice, this list of conditions and the following disclaimer; |
| # redistributions in binary form must reproduce the above copyright |
| # notice, this list of conditions and the following disclaimer in the |
| # documentation and/or other materials provided with the distribution; |
| # neither the name of the copyright holders nor the names of its |
| # contributors may be used to endorse or promote products derived from |
| # this software without specific prior written permission. |
| # |
| # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS |
| # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT |
| # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR |
| # A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT |
| # OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, |
| # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT |
| # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, |
| # DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY |
| # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT |
| # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE |
| # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. |
| |
| """This file defines the ArtifactDB type and some common implementations of |
| ArtifactDB. |
| |
| The database interface defined here does not include any schema information. |
| The database "schema" is defined in the artifact.py file based on the types of |
| artifacts stored in the database. |
| |
| Some common queries can be found in common_queries.py |
| """ |
| |
| from abc import ABC, abstractmethod |
| |
| import copy |
| import json |
| import os |
| from pathlib import Path |
| import shutil |
| from typing import Any, Dict, Iterable, Union, Type, List, Tuple |
| from urllib.parse import urlparse |
| from uuid import UUID |
| |
| try: |
| import gridfs # type: ignore |
| from pymongo import MongoClient # type: ignore |
| |
| MONGO_SUPPORT = True |
| except ModuleNotFoundError: |
| # If pymongo isn't installed, then disable support for it |
| MONGO_SUPPORT = False |
| |
| |
| class ArtifactDB(ABC): |
| """ |
| Abstract base class for all artifact DBs. |
| """ |
| |
| @abstractmethod |
| def __init__(self, uri: str) -> None: |
| """Initialize the database with a URI""" |
| pass |
| |
| @abstractmethod |
| def put(self, key: UUID, artifact: Dict[str, Union[str, UUID]]) -> None: |
| """Insert the artifact into the database with the key""" |
| pass |
| |
| @abstractmethod |
| def upload(self, key: UUID, path: Path) -> None: |
| """Upload the file at path to the database with _id of key""" |
| pass |
| |
| @abstractmethod |
| def __contains__(self, key: Union[UUID, str]) -> bool: |
| """Key can be a UUID or a string. Returns true if item in DB""" |
| pass |
| |
| @abstractmethod |
| def get(self, key: Union[UUID, str]) -> Dict[str, str]: |
| """Key can be a UUID or a string. Returns a dictionary to construct |
| an artifact. |
| """ |
| pass |
| |
| @abstractmethod |
| def downloadFile(self, key: UUID, path: Path) -> None: |
| """Download the file with the _id key to the path. Will overwrite the |
| file if it currently exists.""" |
| pass |
| |
| def searchByName(self, name: str, limit: int) -> Iterable[Dict[str, Any]]: |
| """Returns an iterable of all artifacts in the database that match |
| some name. Note: Not all DB implementations will implement this |
| function""" |
| raise NotImplementedError() |
| |
| def searchByType(self, typ: str, limit: int) -> Iterable[Dict[str, Any]]: |
| """Returns an iterable of all artifacts in the database that match |
| some type. Note: Not all DB implementations will implement this |
| function""" |
| raise NotImplementedError() |
| |
| def searchByNameType( |
| self, name: str, typ: str, limit: int |
| ) -> Iterable[Dict[str, Any]]: |
| """Returns an iterable of all artifacts in the database that match |
| some name and type. Note: Not all DB implementations will implement |
| this function""" |
| raise NotImplementedError() |
| |
| def searchByLikeNameType( |
| self, name: str, typ: str, limit: int |
| ) -> Iterable[Dict[str, Any]]: |
| """Returns an iterable of all artifacts in the database that match |
| some type and a regex name. Note: Not all DB implementations will |
| implement this function""" |
| raise NotImplementedError() |
| |
| |
| class ArtifactMongoDB(ArtifactDB): |
| """ |
| This is a mongodb database connector for storing Artifacts (as defined in |
| artifact.py). |
| |
| This database stores the data in three collections: |
| - artifacts: This stores the json serialized Artifact class |
| - files and chunks: These two collections store the large files required |
| for some artifacts. Within the files collection, the _id is the |
| UUID of the artifact. |
| """ |
| |
| def __init__(self, uri: str) -> None: |
| """Initialize the mongodb connection and grab pointers to the databases |
| uri is the location of the database in a mongodb compatible form. |
| http://dochub.mongodb.org/core/connections. |
| """ |
| # Note: Need "connect=False" so that we don't connect until the first |
| # time we interact with the database. Required for the gem5 running |
| # celery server |
| self.db = MongoClient(host=uri, connect=False).artifact_database |
| self.artifacts = self.db.artifacts |
| self.fs = gridfs.GridFSBucket(self.db, disable_md5=True) |
| |
| def put(self, key: UUID, artifact: Dict[str, Union[str, UUID]]) -> None: |
| """Insert the artifact into the database with the key""" |
| assert artifact["_id"] == key |
| self.artifacts.insert_one(artifact) |
| |
| def upload(self, key: UUID, path: Path) -> None: |
| """Upload the file at path to the database with _id of key""" |
| with open(path, "rb") as f: |
| self.fs.upload_from_stream_with_id(key, str(path), f) |
| |
| def __contains__(self, key: Union[UUID, str]) -> bool: |
| """Key can be a UUID or a string. Returns true if item in DB""" |
| if isinstance(key, UUID): |
| count = self.artifacts.count_documents({"_id": key}, limit=1) |
| else: |
| # This is a hash. Count the number of matches |
| count = self.artifacts.count_documents({"hash": key}, limit=1) |
| |
| return bool(count > 0) |
| |
| def get(self, key: Union[UUID, str]) -> Dict[str, str]: |
| """Key can be a UUID or a string. Returns a dictionary to construct |
| an artifact. |
| """ |
| if isinstance(key, UUID): |
| return self.artifacts.find_one({"_id": key}, limit=1) |
| else: |
| # This is a hash. |
| return self.artifacts.find_one({"hash": key}, limit=1) |
| |
| def downloadFile(self, key: UUID, path: Path) -> None: |
| """Download the file with the _id key to the path. Will overwrite the |
| file if it currently exists.""" |
| with open(path, "wb") as f: |
| self.fs.download_to_stream(key, f) |
| |
| def searchByName(self, name: str, limit: int) -> Iterable[Dict[str, Any]]: |
| """Returns an iterable of all artifacts in the database that match |
| some name.""" |
| for d in self.artifacts.find({"name": name}, limit=limit): |
| yield d |
| |
| def searchByType(self, typ: str, limit: int) -> Iterable[Dict[str, Any]]: |
| """Returns an iterable of all artifacts in the database that match |
| some type.""" |
| for d in self.artifacts.find({"type": typ}, limit=limit): |
| yield d |
| |
| def searchByNameType( |
| self, name: str, typ: str, limit: int |
| ) -> Iterable[Dict[str, Any]]: |
| """Returns an iterable of all artifacts in the database that match |
| some name and type.""" |
| for d in self.artifacts.find({"type": typ, "name": name}, limit=limit): |
| yield d |
| |
| def searchByLikeNameType( |
| self, name: str, typ: str, limit: int |
| ) -> Iterable[Dict[str, Any]]: |
| """Returns an iterable of all artifacts in the database that match |
| some type and a regex name.""" |
| |
| data = self.artifacts.find( |
| {"type": typ, "name": {"$regex": "{}".format(name)}}, limit=limit |
| ) |
| for d in data: |
| yield d |
| |
| |
| class ArtifactFileDB(ArtifactDB): |
| """ |
| This is a file-based database where Artifacts (as defined in artifacts.py) |
| are stored in a JSON file. |
| |
| This database stores a list of serialized artifacts in a JSON file. |
| This database is not thread-safe. |
| |
| If the user specifies a valid path in the environment variable |
| GEM5ART_STORAGE then this database will copy all artifacts to that |
| directory named with their UUIDs. |
| """ |
| |
| class ArtifactEncoder(json.JSONEncoder): |
| def default(self, obj): |
| if isinstance(obj, UUID): |
| return str(obj) |
| return ArtifactFileDB.ArtifactEncoder(self, obj) |
| |
| _json_file: Path |
| _uuid_artifact_map: Dict[str, Dict[str, str]] |
| _hash_uuid_map: Dict[str, List[str]] |
| _storage_enabled: bool |
| _storage_path: Path |
| |
| def __init__(self, uri: str) -> None: |
| """Initialize the file-driven database from a JSON file. |
| If the file doesn't exist, a new file will be created. |
| """ |
| parsed_uri = urlparse(uri) |
| # using urlparse to parse relative/absolute file path |
| # abs path: urlparse("file:///path/to/file") -> |
| # (netloc='', path='/path/to/file') |
| # rel path: urlparse("file://path/to/file") -> |
| # (netloc='path', path='/to/file') |
| # so, the filepath would be netloc+path for both cases |
| self._json_file = Path(parsed_uri.netloc) / Path(parsed_uri.path) |
| storage_path = os.environ.get("GEM5ART_STORAGE", "") |
| self._storage_enabled = True if storage_path else False |
| self._storage_path = Path(storage_path) |
| if ( |
| self._storage_enabled |
| and self._storage_path.exists() |
| and not self._storage_path.is_dir() |
| ): |
| raise Exception( |
| f"GEM5ART_STORAGE={storage_path} exists and is not a directory" |
| ) |
| if self._storage_enabled: |
| os.makedirs(self._storage_path, exist_ok=True) |
| |
| self._uuid_artifact_map, self._hash_uuid_map = self._load_from_file( |
| self._json_file |
| ) |
| |
| def put(self, key: UUID, artifact: Dict[str, Union[str, UUID]]) -> None: |
| """Insert the artifact into the database with the key.""" |
| assert artifact["_id"] == key |
| assert isinstance(artifact["hash"], str) |
| self.insert_artifact(key, artifact["hash"], artifact) |
| |
| def upload(self, key: UUID, path: Path) -> None: |
| """Copy the artifact to the folder specified by GEM5ART_STORAGE.""" |
| if not self._storage_enabled: |
| return |
| src_path = path |
| dst_path = self._storage_path / str(key) |
| if not dst_path.exists(): |
| shutil.copy2(src_path, dst_path) |
| |
| def __contains__(self, key: Union[UUID, str]) -> bool: |
| """Key can be a UUID or a string. Returns true if item in DB""" |
| if isinstance(key, UUID): |
| return self.has_uuid(key) |
| return self.has_hash(key) |
| |
| def get(self, key: Union[UUID, str]) -> Dict[str, str]: |
| """Key can be a UUID or a string. Returns a dictionary to construct |
| an artifact. |
| """ |
| artifact: List[Dict[str, str]] = [] |
| if isinstance(key, UUID): |
| artifact = list(self.get_artifact_by_uuid(key)) |
| else: |
| # This is a hash. |
| artifact = list(self.get_artifact_by_hash(key)) |
| return artifact[0] |
| |
| def downloadFile(self, key: UUID, path: Path) -> None: |
| """Copy the file from the storage to specified path.""" |
| assert path.exists() |
| if not self._storage_enabled: |
| return |
| src_path = self._storage_path / str(key) |
| dst_path = path |
| shutil.copy2(src_path, dst_path) |
| |
| def _load_from_file( |
| self, json_file: Path |
| ) -> Tuple[Dict[str, Dict[str, str]], Dict[str, List[str]]]: |
| uuid_mapping: Dict[str, Dict[str, str]] = {} |
| hash_mapping: Dict[str, List[str]] = {} |
| if json_file.exists(): |
| with open(json_file, "r") as f: |
| j = json.load(f) |
| for an_artifact in j: |
| the_uuid = an_artifact["_id"] |
| the_hash = an_artifact["hash"] |
| uuid_mapping[the_uuid] = an_artifact |
| if not the_hash in hash_mapping: |
| hash_mapping[the_hash] = [] |
| hash_mapping[the_hash].append(the_uuid) |
| return uuid_mapping, hash_mapping |
| |
| def _save_to_file(self, json_file: Path) -> None: |
| content = list(self._uuid_artifact_map.values()) |
| with open(json_file, "w") as f: |
| json.dump(content, f, indent=4, cls=ArtifactFileDB.ArtifactEncoder) |
| |
| def has_uuid(self, the_uuid: UUID) -> bool: |
| return str(the_uuid) in self._uuid_artifact_map |
| |
| def has_hash(self, the_hash: str) -> bool: |
| return the_hash in self._hash_uuid_map |
| |
| def get_artifact_by_uuid(self, the_uuid: UUID) -> Iterable[Dict[str, str]]: |
| uuid_str = str(the_uuid) |
| if not uuid_str in self._uuid_artifact_map: |
| return |
| yield self._uuid_artifact_map[uuid_str] |
| |
| def get_artifact_by_hash(self, the_hash: str) -> Iterable[Dict[str, str]]: |
| if not the_hash in self._hash_uuid_map: |
| return |
| for the_uuid in self._hash_uuid_map[the_hash]: |
| yield self._uuid_artifact_map[the_uuid] |
| |
| def insert_artifact( |
| self, |
| the_uuid: UUID, |
| the_hash: str, |
| the_artifact: Dict[str, Union[str, UUID]], |
| ) -> bool: |
| """ |
| Put the artifact to the database. |
| |
| Return True if the artifact uuid does not exist in the database prior |
| to calling this function; return False otherwise. |
| """ |
| uuid_str = str(the_uuid) |
| if uuid_str in self._uuid_artifact_map: |
| return False |
| artifact_copy = copy.deepcopy(the_artifact) |
| artifact_copy["_id"] = str(artifact_copy["_id"]) |
| self._uuid_artifact_map[uuid_str] = artifact_copy # type: ignore |
| if not the_hash in self._hash_uuid_map: |
| self._hash_uuid_map[the_hash] = [] |
| self._hash_uuid_map[the_hash].append(uuid_str) |
| self._save_to_file(self._json_file) |
| return True |
| |
| def find_exact( |
| self, attr: Dict[str, str], limit: int |
| ) -> Iterable[Dict[str, Any]]: |
| """ |
| Return all artifacts such that, for every yielded artifact, |
| and for every (k,v) in attr, the attribute `k` of the artifact has |
| the value of `v`. |
| """ |
| count = 0 |
| if count >= limit: |
| return |
| for artifact in self._uuid_artifact_map.values(): |
| #https://docs.python.org/3/library/stdtypes.html#frozenset.issubset |
| if attr.items() <= artifact.items(): |
| yield artifact |
| |
| |
| _db = None |
| |
| if MONGO_SUPPORT: |
| _default_uri = "mongodb://localhost:27017" |
| else: |
| _default_uri = "file://db.json" |
| |
| _db_schemes: Dict[str, Type[ArtifactDB]] = {"file": ArtifactFileDB} |
| if MONGO_SUPPORT: |
| _db_schemes["mongodb"] = ArtifactMongoDB |
| |
| |
| def _getDBType(uri: str) -> Type[ArtifactDB]: |
| """Internal function to take a URI and return a class that can be |
| constructed with that URI. For instance "mongodb://localhost" will return |
| an ArtifactMongoDB. More types will be added in the future. |
| |
| Supported types: |
| **ArtifactMongoDB**: mongodb://... |
| See http://dochub.mongodb.org/core/connections for details. |
| **ArtifactFileDB**: file://... |
| A simple flat file database with optional storage for the binary |
| artifacts. The filepath is where the json file is stored and the |
| data storage can be specified with GEM5ART_STORAGE |
| """ |
| result = urlparse(uri) |
| if result.scheme in _db_schemes: |
| return _db_schemes[result.scheme] |
| else: |
| raise Exception(f"Cannot find DB type for {uri}") |
| |
| |
| def getDBConnection(uri: str = "") -> ArtifactDB: |
| """Returns the database connection |
| |
| uri: a string representing the URI of the database. See _getDBType for |
| details. If no URI is given we use the default |
| (mongodb://localhost:27017) or the value in the GEM5ART_DB environment |
| variable. |
| |
| If the connection has not been established, this will create a new |
| connection. If the connection has been established, this will replace the |
| connection if the uri input is non-empy. |
| """ |
| global _db |
| |
| # mypy bug: https://github.com/python/mypy/issues/5423 |
| if _db is not None and not uri: # type: ignore[unreachable] |
| # If we have already established a connection, use that |
| return _db # type: ignore[unreachable] |
| |
| if not uri: |
| uri = os.environ.get("GEM5ART_DB", _default_uri) |
| |
| typ = _getDBType(uri) |
| _db = typ(uri) |
| |
| return _db |
