diff --git a/swh/provenance/cli.py b/swh/provenance/cli.py --- a/swh/provenance/cli.py +++ b/swh/provenance/cli.py @@ -161,6 +161,7 @@ lower=reuse, mindepth=min_depth, ) + provenance.close() def generate_revision_tuples( @@ -192,6 +193,7 @@ for origin in origins: origin_add(provenance, archive, [origin]) + provenance.close() def generate_origin_tuples(filename: str) -> Generator[Tuple[str, bytes], None, None]: @@ -220,6 +222,7 @@ ) else: print(f"Cannot find a content with the id {swhid}") + provenance.close() @cli.command(name="find-all") @@ -239,3 +242,4 @@ f"{occur.origin}, " f"{os.fsdecode(occur.path)}" ) + provenance.close() diff --git a/swh/provenance/interface.py b/swh/provenance/interface.py --- a/swh/provenance/interface.py +++ b/swh/provenance/interface.py @@ -65,6 +65,11 @@ @runtime_checkable class ProvenanceStorageInterface(Protocol): + @remote_api_endpoint("close") + def close(self) -> None: + """Close connection to the storage and release resources.""" + ... + @remote_api_endpoint("content_add") def content_add( self, cnts: Union[Iterable[Sha1Git], Dict[Sha1Git, Optional[datetime]]] @@ -198,6 +203,10 @@ class ProvenanceInterface(Protocol): storage: ProvenanceStorageInterface + def close(self) -> None: + """Close connection to the underlying `storage` and release resources.""" + ... + def flush(self) -> None: """Flush internal cache to the underlying `storage`.""" ... diff --git a/swh/provenance/mongo/backend.py b/swh/provenance/mongo/backend.py --- a/swh/provenance/mongo/backend.py +++ b/swh/provenance/mongo/backend.py @@ -25,6 +25,9 @@ def __init__(self, db: pymongo.database.Database): self.db = db + def close(self) -> None: + pass + def content_add( self, cnts: Union[Iterable[Sha1Git], Dict[Sha1Git, Optional[datetime]]] ) -> bool: diff --git a/swh/provenance/postgresql/provenance.py b/swh/provenance/postgresql/provenance.py --- a/swh/provenance/postgresql/provenance.py +++ b/swh/provenance/postgresql/provenance.py @@ -60,6 +60,9 @@ def denormalized(self) -> bool: return "denormalized" in self.flavor + def close(self) -> None: + pass + def content_add( self, cnts: Union[Iterable[Sha1Git], Dict[Sha1Git, Optional[datetime]]] ) -> bool: diff --git a/swh/provenance/provenance.py b/swh/provenance/provenance.py --- a/swh/provenance/provenance.py +++ b/swh/provenance/provenance.py @@ -77,6 +77,9 @@ def clear_caches(self) -> None: self.cache = new_cache() + def close(self) -> None: + self.storage.close() + def flush(self) -> None: # Revision-content layer insertions ############################################ diff --git a/swh/provenance/tests/conftest.py b/swh/provenance/tests/conftest.py --- a/swh/provenance/tests/conftest.py +++ b/swh/provenance/tests/conftest.py @@ -5,7 +5,7 @@ from datetime import datetime, timedelta, timezone from os import path -from typing import Any, Dict, Iterable, Iterator +from typing import Any, Dict, Generator, Iterable, Iterator from _pytest.fixtures import SubRequest import msgpack @@ -70,23 +70,28 @@ provenance_postgresqldb: Dict[str, str], mongodb: pymongo.database.Database, swh_rpc_client: ProvenanceStorageRPCClient, -) -> ProvenanceStorageInterface: +) -> Generator[ProvenanceStorageInterface, None, None]: """Return a working and initialized ProvenanceStorageInterface object""" if request.param == "rpcapi": assert isinstance(swh_rpc_client, ProvenanceStorageInterface) - return swh_rpc_client + yield swh_rpc_client + swh_rpc_client.close() elif request.param == "mongodb": from swh.provenance.mongo.backend import ProvenanceStorageMongoDb - return ProvenanceStorageMongoDb(mongodb) + mongodb_storage = ProvenanceStorageMongoDb(mongodb) + yield mongodb_storage + mongodb_storage.close() else: # in test sessions, we DO want to raise any exception occurring at commit time - return get_provenance_storage( + storage = get_provenance_storage( cls=request.param, db=provenance_postgresqldb, raise_on_commit=True ) + yield storage + storage.close() provenance_postgresql = postgresql("postgresql_proc", dbname="provenance_tests") @@ -95,7 +100,7 @@ @pytest.fixture def provenance( provenance_postgresql: psycopg2.extensions.connection, -) -> ProvenanceInterface: +) -> Generator[ProvenanceInterface, None, None]: """Return a working and initialized ProvenanceInterface object""" from swh.core.cli.db import populate_database_for_package @@ -104,11 +109,13 @@ "swh.provenance", provenance_postgresql.dsn, flavor="with-path" ) # in test sessions, we DO want to raise any exception occurring at commit time - return get_provenance( + provenance = get_provenance( cls="postgresql", db=provenance_postgresql.get_dsn_parameters(), raise_on_commit=True, ) + yield provenance + provenance.close() @pytest.fixture