diff --git a/nucliadb/src/nucliadb/common/maindb/pg.py b/nucliadb/src/nucliadb/common/maindb/pg.py index a8e5b7f3e7..ac67a5965c 100644 --- a/nucliadb/src/nucliadb/common/maindb/pg.py +++ b/nucliadb/src/nucliadb/common/maindb/pg.py @@ -275,9 +275,10 @@ async def initialize(self): async def finalize(self): async with self._lock: - await self.pool.close() - self.initialized = False - self.metrics_task.cancel() + if self.initialized: + await self.pool.close() + self.initialized = False + self.metrics_task.cancel() async def _report_metrics_task(self): while True: diff --git a/nucliadb/tests/ndbfixtures/maindb.py b/nucliadb/tests/ndbfixtures/maindb.py index a84418e6cc..a90d358bcb 100644 --- a/nucliadb/tests/ndbfixtures/maindb.py +++ b/nucliadb/tests/ndbfixtures/maindb.py @@ -106,13 +106,21 @@ async def maindb_driver(request: FixtureRequest) -> AsyncIterator[Driver]: async def pg_maindb_settings(pg): url = f"postgresql://postgres:postgres@{pg[0]}:{pg[1]}/postgres" + # We want to be sure schema migrations are always run. As some tests use + # this fixture and create their own driver, we need to create one here and + # run the migrations, so the maindb_settings fixture can still be generic + # and pg migrations are run driver = PGDriver(url=url, connection_pool_min_size=2, connection_pool_max_size=2) await driver.initialize() await run_pg_schema_migrations(driver) + await driver.finalize() - return DriverSettings( + yield DriverSettings( driver=DriverConfig.PG, driver_pg_url=url, + driver_pg_connection_pool_min_size=10, + driver_pg_connection_pool_max_size=10, + driver_pg_connection_pool_acquire_timeout_ms=200, ) @@ -129,7 +137,12 @@ async def pg_maindb_driver(pg_maindb_settings: DriverSettings): await cur.execute("DROP table IF EXISTS resources") await cur.execute("DROP table IF EXISTS catalog") - driver = PGDriver(url=url) + driver = PGDriver( + url=url, + connection_pool_min_size=pg_maindb_settings.driver_pg_connection_pool_min_size, + connection_pool_max_size=pg_maindb_settings.driver_pg_connection_pool_max_size, + acquire_timeout_ms=pg_maindb_settings.driver_pg_connection_pool_acquire_timeout_ms, + ) await driver.initialize() await run_pg_schema_migrations(driver) diff --git a/nucliadb/tests/ndbfixtures/utils.py b/nucliadb/tests/ndbfixtures/utils.py index 8eb5abd555..f360344d36 100644 --- a/nucliadb/tests/ndbfixtures/utils.py +++ b/nucliadb/tests/ndbfixtures/utils.py @@ -22,7 +22,7 @@ from typing import Callable, Optional from fastapi import FastAPI -from httpx import AsyncClient +from httpx import ASGITransport, AsyncClient from nucliadb.search import API_PREFIX @@ -41,7 +41,8 @@ def _make_client_fixture( if root is False: client_base_url = f"{client_base_url}/{API_PREFIX}/v{version}" - client = AsyncClient(app=application, base_url=client_base_url) + transport = ASGITransport(app=application) # type: ignore + client = AsyncClient(transport=transport, base_url=client_base_url) client.headers["X-NUCLIADB-ROLES"] = ";".join([role.value for role in roles]) client.headers["X-NUCLIADB-USER"] = user diff --git a/nucliadb/tests/reader/benchmarks/test_benchmarks_resource.py b/nucliadb/tests/reader/benchmarks/test_benchmarks_resource.py index 1db0c8592d..66b035290f 100644 --- a/nucliadb/tests/reader/benchmarks/test_benchmarks_resource.py +++ b/nucliadb/tests/reader/benchmarks/test_benchmarks_resource.py @@ -23,6 +23,7 @@ from httpx import AsyncClient from nucliadb.ingest.orm.resource import Resource +from nucliadb.reader.api.v1.router import KB_PREFIX, RESOURCE_PREFIX from nucliadb_utils.tests.asyncbenchmark import AsyncBenchmarkFixture @@ -38,36 +39,32 @@ @pytest.mark.deploy_modes("component") async def test_get_resource_all( nucliadb_reader: AsyncClient, - test_resource: Resource, + full_resource: Resource, asyncbenchmark: AsyncBenchmarkFixture, ) -> None: - # resource = full_resource - # kbid = resource.kb.kbid - # rid = resource.uuid + resource = full_resource + kbid = resource.kb.kbid + rid = resource.uuid - # resp = await asyncbenchmark( - # nucliadb_reader.get, - # f"/{KB_PREFIX}/{kbid}/{RESOURCE_PREFIX}/{rid}", - # params={ - # "show": ["basic", "origin", "relations", "values", "extracted"], - # "field_type": [ - # "text", - # "link", - # "file", - # "conversation", - # ], - # "extracted": [ - # "metadata", - # "vectors", - # "large_metadata", - # "text", - # "link", - # "file", - # ], - # }, - # ) - # assert resp.status_code == 200 - async def function(a): - return a - - assert await asyncbenchmark(function, 10) == 10 + resp = await asyncbenchmark( + nucliadb_reader.get, + f"/{KB_PREFIX}/{kbid}/{RESOURCE_PREFIX}/{rid}", + params={ + "show": ["basic", "origin", "relations", "values", "extracted"], + "field_type": [ + "text", + "link", + "file", + "conversation", + ], + "extracted": [ + "metadata", + "vectors", + "large_metadata", + "text", + "link", + "file", + ], + }, + ) + assert resp.status_code == 200 diff --git a/nucliadb_utils/src/nucliadb_utils/audit/basic.py b/nucliadb_utils/src/nucliadb_utils/audit/basic.py index ac102abbb8..8096adbeb0 100644 --- a/nucliadb_utils/src/nucliadb_utils/audit/basic.py +++ b/nucliadb_utils/src/nucliadb_utils/audit/basic.py @@ -54,7 +54,7 @@ async def report_and_send( ): logger.debug(f"AUDIT {audit_type} {kbid} {user} {origin} {rid} {audit_fields}") - async def visited( + def visited( self, kbid: str, uuid: str,