nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2798__py3-none-any.whl

nucliadb/common/maindb/pg.py

@@ -20,65 +20,108 @@
 from __future__ import annotations
 
 import asyncio
-from typing import Any, AsyncGenerator, Optional, Union
+import logging
+from contextlib import asynccontextmanager
+from typing import Any, AsyncGenerator, Optional
 
-import asyncpg
 import backoff
+import psycopg
+import psycopg_pool
 
 from nucliadb.common.maindb.driver import DEFAULT_SCAN_LIMIT, Driver, Transaction
+from nucliadb.common.maindb.exceptions import ConflictError
+from nucliadb_telemetry import metrics
 
 RETRIABLE_EXCEPTIONS = (
-    asyncpg.CannotConnectNowError,
+    psycopg_pool.PoolTimeout,
     OSError,
     ConnectionResetError,
 )
 
-CREATE_TABLE = """
-CREATE TABLE IF NOT EXISTS resources (
-    key TEXT PRIMARY KEY,
-    value BYTEA
-);
-"""
+logger = logging.getLogger(__name__)
+
+# Request Metrics
+pg_observer = metrics.Observer(
+    "pg_client",
+    labels={"type": ""},
+)
+
+# Pool metrics
+POOL_METRICS_COUNTERS = {
+    # Requests for a connection to the pool
+    "requests_num": metrics.Counter("pg_client_pool_requests_total"),
+    "requests_queued": metrics.Counter("pg_client_pool_requests_queued_total"),
+    "requests_errors": metrics.Counter("pg_client_pool_requests_errors_total"),
+    "requests_wait_ms": metrics.Counter("pg_client_pool_requests_queued_seconds_total"),
+    "usage_ms": metrics.Counter("pg_client_pool_requests_usage_seconds_total"),
+    # Pool opening a connection to PG
+    "connections_num": metrics.Counter("pg_client_pool_connections_total"),
+    "connections_ms": metrics.Counter("pg_client_pool_connections_seconds_total"),
+}
+POOL_METRICS_GAUGES = {
+    "pool_size": metrics.Gauge("pg_client_pool_connections_open"),
+    # The two below most likely change too rapidly to be useful in a metric
+    "pool_available": metrics.Gauge("pg_client_pool_connections_available"),
+    "requests_waiting": metrics.Gauge("pg_client_pool_requests_waiting"),
+}
 
 
 class DataLayer:
-    def __init__(self, connection: Union[asyncpg.Connection, asyncpg.Pool]):
+    def __init__(self, connection: psycopg.AsyncConnection):
         self.connection = connection
-        self.lock = asyncio.Lock()
 
-    async def get(self, key: str) -> Optional[bytes]:
-        async with self.lock:
-            return await self.connection.fetchval(
-                "SELECT value FROM resources WHERE key = $1", key
-            )
+    async def get(self, key: str, select_for_update: bool = False) -> Optional[bytes]:
+        with pg_observer({"type": "get"}):
+            statement = "SELECT value FROM resources WHERE key = %s"
+            if select_for_update:
+                statement += " FOR UPDATE"
+            async with self.connection.cursor() as cur:
+                await cur.execute(statement, (key,))
+                row = await cur.fetchone()
+                return row[0] if row else None
 
     async def set(self, key: str, value: bytes) -> None:
-        async with self.lock:
-            await self.connection.execute(
-                """
-INSERT INTO resources (key, value)
-VALUES ($1, $2)
-ON CONFLICT (key)
-DO UPDATE SET value = EXCLUDED.value
-""",
-                key,
-                value,
-            )
+        with pg_observer({"type": "set"}):
+            async with self.connection.cursor() as cur:
+                await cur.execute(
+                    "INSERT INTO resources (key, value) "
+                    "VALUES (%s, %s) "
+                    "ON CONFLICT (key) "
+                    "DO UPDATE SET value = EXCLUDED.value",
+                    (key, value),
+                )
+
+    async def insert(self, key: str, value: bytes) -> None:
+        with pg_observer({"type": "insert"}):
+            async with self.connection.cursor() as cur:
+                try:
+                    await cur.execute(
+                        "INSERT INTO resources (key, value) VALUES (%s, %s) ",
+                        (key, value),
+                    )
+                except psycopg.errors.UniqueViolation:
+                    raise ConflictError(key)
 
     async def delete(self, key: str) -> None:
-        async with self.lock:
-            await self.connection.execute("DELETE FROM resources WHERE key = $1", key)
-
-    async def batch_get(self, keys: list[str]) -> list[Optional[bytes]]:
-        async with self.lock:
-            records = {
-                record["key"]: record["value"]
-                for record in await self.connection.fetch(
-                    "SELECT key, value FROM resources WHERE key = ANY($1)", keys
-                )
-            }
-        # get sorted by keys
-        return [records.get(key) for key in keys]
+        with pg_observer({"type": "delete"}):
+            async with self.connection.cursor() as cur:
+                await cur.execute("DELETE FROM resources WHERE key = %s", (key,))
+
+    async def delete_by_prefix(self, prefix: str) -> None:
+        with pg_observer({"type": "delete_by_prefix"}):
+            async with self.connection.cursor() as cur:
+                await cur.execute("DELETE FROM resources WHERE key LIKE %s", (prefix + "%",))
+
+    async def batch_get(self, keys: list[str], select_for_update: bool = False) -> list[Optional[bytes]]:
+        with pg_observer({"type": "batch_get"}):
+            async with self.connection.cursor() as cur:
+                statement = "SELECT key, value FROM resources WHERE key = ANY(%s)"
+                if select_for_update:
+                    statement += " FOR UPDATE"
+                await cur.execute(statement, (keys,))
+                records = {record[0]: record[1] for record in await cur.fetchall()}
+            # get sorted by keys
+            return [records.get(key) for key in keys]
 
     async def scan_keys(
         self,
@@ -86,26 +129,25 @@ DO UPDATE SET value = EXCLUDED.value
         limit: int = DEFAULT_SCAN_LIMIT,
         include_start: bool = True,
     ) -> AsyncGenerator[str, None]:
-        query = """SELECT key FROM resources
-WHERE key LIKE $1
-ORDER BY key
-"""
+        query = "SELECT key FROM resources WHERE key LIKE %s ORDER BY key"
+
         args: list[Any] = [prefix + "%"]
         if limit > 0:
-            query += " LIMIT $2"
+            query += " LIMIT %s"
             args.append(limit)
-        async with self.lock:
-            async for record in self.connection.cursor(query, *args):
-                if not include_start and record["key"] == prefix:
-                    continue
-                yield record["key"]
+        with pg_observer({"type": "scan_keys"}):
+            async with self.connection.cursor() as cur:
+                async for record in cur.stream(query, args):
+                    if not include_start and record[0] == prefix:
+                        continue
+                    yield record[0]
 
     async def count(self, match: str) -> int:
-        async with self.lock:
-            results = await self.connection.fetch(
-                "SELECT count(*) FROM resources WHERE key LIKE $1", match + "%"
-            )
-        return results[0]["count"]
+        with pg_observer({"type": "count"}):
+            async with self.connection.cursor() as cur:
+                await cur.execute("SELECT count(*) FROM resources WHERE key LIKE %s", (match + "%",))
+                row = await cur.fetchone()
+                return row[0]  # type: ignore
 
 
 class PGTransaction(Transaction):
@@ -113,62 +155,60 @@ class PGTransaction(Transaction):
 
     def __init__(
         self,
-        pool: asyncpg.Pool,
-        connection: asyncpg.Connection,
-        txn: Any,
         driver: PGDriver,
+        connection: psycopg.AsyncConnection,
     ):
-        self.pool = pool
+        self.driver = driver
         self.connection = connection
         self.data_layer = DataLayer(connection)
-        self.txn = txn
-        self.driver = driver
         self.open = True
-        self._lock = asyncio.Lock()
 
     async def abort(self):
-        async with self._lock:
+        with pg_observer({"type": "rollback"}):
             if self.open:
                 try:
-                    await self.txn.rollback()
+                    await self.connection.rollback()
                 finally:
                     self.open = False
-                    await self.connection.close()
 
     async def commit(self):
-        async with self._lock:
+        with pg_observer({"type": "commit"}):
             try:
-                await self.txn.commit()
+                await self.connection.commit()
             except Exception:
-                await self.txn.rollback()
+                await self.connection.rollback()
                 raise
             finally:
                 self.open = False
-                await self.connection.close()
 
-    async def batch_get(self, keys: list[str]):
-        return await self.data_layer.batch_get(keys)
+    async def batch_get(self, keys: list[str], for_update: bool = True):
+        return await self.data_layer.batch_get(keys, select_for_update=for_update)
 
-    async def get(self, key: str) -> Optional[bytes]:
-        return await self.data_layer.get(key)
+    async def get(self, key: str, for_update: bool = True) -> Optional[bytes]:
+        return await self.data_layer.get(key, select_for_update=for_update)
 
     async def set(self, key: str, value: bytes):
         await self.data_layer.set(key, value)
 
+    async def insert(self, key: str, value: bytes):
+        await self.data_layer.insert(key, value)
+
     async def delete(self, key: str):
         await self.data_layer.delete(key)
 
-    @backoff.on_exception(
-        backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=2
-    )
+    async def delete_by_prefix(self, prefix: str) -> None:
+        await self.data_layer.delete_by_prefix(prefix)
+
+    @backoff.on_exception(backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=2)
     async def keys(
         self,
         match: str,
         count: int = DEFAULT_SCAN_LIMIT,
         include_start: bool = True,
     ):
-        async with self.pool.acquire() as conn, conn.transaction():
-            # all txn implementations implement this API outside of the current txn
+        # Check out a new connection to guarantee that the cursor iteration does not
+        # run concurrently with other queries
+        async with self.driver._get_connection() as conn, conn.transaction():
             dl = DataLayer(conn)
             async for key in dl.scan_keys(match, count, include_start=include_start):
                 yield key
@@ -180,8 +220,7 @@ class PGTransaction(Transaction):
 class ReadOnlyPGTransaction(Transaction):
     driver: PGDriver
 
-    def __init__(self, pool: asyncpg.Pool, driver: PGDriver):
-        self.pool = pool
+    def __init__(self, driver: PGDriver):
         self.driver = driver
         self.open = True
 
@@ -192,11 +231,15 @@ class ReadOnlyPGTransaction(Transaction):
     async def commit(self):
         raise Exception("Cannot commit transaction in read only mode")
 
-    async def batch_get(self, keys: list[str]):
-        return await DataLayer(self.pool).batch_get(keys)
+    @backoff.on_exception(backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=3)
+    async def batch_get(self, keys: list[str], for_update: bool = False):
+        async with self.driver._get_connection() as conn:
+            return await DataLayer(conn).batch_get(keys, select_for_update=False)
 
-    async def get(self, key: str) -> Optional[bytes]:
-        return await DataLayer(self.pool).get(key)
+    @backoff.on_exception(backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=3)
+    async def get(self, key: str, for_update: bool = False) -> Optional[bytes]:
+        async with self.driver._get_connection() as conn:
+            return await DataLayer(conn).get(key, select_for_update=False)
 
     async def set(self, key: str, value: bytes):
         raise Exception("Cannot set in read only transaction")
@@ -204,62 +247,114 @@ class ReadOnlyPGTransaction(Transaction):
     async def delete(self, key: str):
         raise Exception("Cannot delete in read only transaction")
 
+    @backoff.on_exception(backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=3)
     async def keys(
         self,
         match: str,
         count: int = DEFAULT_SCAN_LIMIT,
         include_start: bool = True,
     ):
-        async with self.pool.acquire() as conn, conn.transaction():
-            # all txn implementations implement this API outside of the current txn
+        async with self.driver._get_connection() as conn, conn.transaction():
             dl = DataLayer(conn)
             async for key in dl.scan_keys(match, count, include_start=include_start):
                 yield key
 
+    @backoff.on_exception(backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=3)
     async def count(self, match: str) -> int:
-        return await DataLayer(self.pool).count(match)
+        async with self.driver._get_connection() as conn:
+            return await DataLayer(conn).count(match)
 
 
 class PGDriver(Driver):
-    pool: asyncpg.Pool
+    pool: psycopg_pool.AsyncConnectionPool
 
-    def __init__(self, url: str, connection_pool_max_size: int = 10):
+    def __init__(
+        self,
+        url: str,
+        connection_pool_min_size: int = 10,
+        connection_pool_max_size: int = 10,
+        acquire_timeout_ms: int = 200,
+    ):
         self.url = url
+        self.connection_pool_min_size = connection_pool_min_size
         self.connection_pool_max_size = connection_pool_max_size
+        self.acquire_timeout_ms = acquire_timeout_ms
         self._lock = asyncio.Lock()
 
     async def initialize(self):
         async with self._lock:
             if self.initialized is False:
-                self.pool = await asyncpg.create_pool(
+                self.pool = psycopg_pool.AsyncConnectionPool(
                     self.url,
+                    min_size=self.connection_pool_min_size,
                     max_size=self.connection_pool_max_size,
+                    check=psycopg_pool.AsyncConnectionPool.check_connection,
+                    open=False,
                 )
-
-                # check if table exists
-                try:
-                    async with self.pool.acquire() as conn:
-                        await conn.execute(CREATE_TABLE)
-                except asyncpg.exceptions.UniqueViolationError:  # pragma: no cover
-                    pass
+                await self.pool.open()
 
             self.initialized = True
+            self.metrics_task = asyncio.create_task(self._report_metrics_task())
 
     async def finalize(self):
         async with self._lock:
-            await self.pool.close()
-            self.initialized = False
-
-    @backoff.on_exception(
-        backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=3
-    )
-    async def begin(
-        self, read_only: bool = False
-    ) -> Union[PGTransaction, ReadOnlyPGTransaction]:
+            if self.initialized:
+                await self.pool.close()
+                self.initialized = False
+                self.metrics_task.cancel()
+
+    async def _report_metrics_task(self):
+        while True:
+            self._report_metrics()
+            await asyncio.sleep(60)
+
+    def _report_metrics(self):
+        if not self.initialized:
+            return
+
+        metrics = self.pool.pop_stats()
+        for key, metric in POOL_METRICS_COUNTERS.items():
+            value = metrics.get(key, 0)
+            if key.endswith("_ms"):
+                value /= 1000
+            metric.inc(value=value)
+
+        for key, metric in POOL_METRICS_GAUGES.items():
+            value = metrics.get(key, 0)
+            metric.set(value)
+
+    @asynccontextmanager
+    async def transaction(self, read_only: bool = False) -> AsyncGenerator[Transaction, None]:
         if read_only:
-            return ReadOnlyPGTransaction(self.pool, driver=self)
+            yield ReadOnlyPGTransaction(self)
         else:
-            conn: asyncpg.Connection = await self.pool.acquire()
-            txn = conn.transaction()
-            await txn.start()
-            return PGTransaction(self.pool, conn, txn, driver=self)
+            async with self._get_connection() as conn:
+                txn = PGTransaction(self, conn)
+                try:
+                    yield txn
+                finally:
+                    if txn.open:
+                        await txn.abort()
+
+    @asynccontextmanager
+    async def _get_connection(self) -> AsyncGenerator[psycopg.AsyncConnection, None]:
+        timeout = self.acquire_timeout_ms / 1000
+        # Manual retry loop since backoff.on_exception does not play well with async context managers
+        retries = 0
+        while True:
+            with pg_observer({"type": "acquire"}):
+                try:
+                    async with self.pool.connection(timeout=timeout) as conn:
+                        yield conn
+                        return
+                except psycopg_pool.PoolTimeout as e:
+                    logger.warning(
+                        f"Timeout getting connection from the pool, backing off. Retries = {retries}"
+                    )
+                    if retries < 3:
+                        await asyncio.sleep(1)
+                        retries += 1
+                    else:
+                        raise e
+                except Exception as e:
+                    raise e

nucliadb/common/maindb/utils.py

@@ -16,6 +16,8 @@
 #
 # You should have received a copy of the GNU Affero General Public License
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
+from warnings import warn
+
 from nucliadb.common.maindb.driver import Driver
 from nucliadb.common.maindb.exceptions import UnsetUtility
 from nucliadb.ingest.settings import DriverConfig, settings
@@ -23,19 +25,11 @@ from nucliadb_utils.exceptions import ConfigurationError
 from nucliadb_utils.utilities import Utility, clean_utility, get_utility, set_utility
 
 try:
-    from nucliadb.common.maindb.redis import RedisDriver
-
-    REDIS = True
-except ImportError:  # pragma: no cover
-    REDIS = False
-
-try:
-    from nucliadb.common.maindb.tikv import TiKVDriver
+    from nucliadb.common.maindb.local import LocalDriver
 
-    TIKV = True
+    FILES = True
 except ImportError:  # pragma: no cover
-    TIKV = False
-
+    FILES = False
 
 try:
     from nucliadb.common.maindb.pg import PGDriver
@@ -44,13 +38,6 @@ try:
 except ImportError:  # pragma: no cover
     PG = False
 
-try:
-    from nucliadb.common.maindb.local import LocalDriver
-
-    FILES = True
-except ImportError:  # pragma: no cover
-    FILES = False
-
 
 def get_driver() -> Driver:
     driver = get_utility(Utility.MAINDB_DRIVER)
@@ -64,32 +51,16 @@ async def setup_driver() -> Driver:
     if driver is not None:
         return driver
 
-    if settings.driver == DriverConfig.REDIS:
-        if not REDIS:
-            raise ConfigurationError("`redis` python package not installed.")
-        if settings.driver_redis_url is None:
-            raise ConfigurationError("No DRIVER_REDIS_URL env var defined.")
-
-        redis_driver = RedisDriver(settings.driver_redis_url)
-        set_utility(Utility.MAINDB_DRIVER, redis_driver)
-    elif settings.driver == DriverConfig.TIKV:
-        if not TIKV:
-            raise ConfigurationError("`tikv_client` python package not installed.")
-        if settings.driver_tikv_url is None:
-            raise ConfigurationError("No DRIVER_TIKV_URL env var defined.")
-
-        tikv_driver = TiKVDriver(
-            settings.driver_tikv_url, settings.driver_tikv_connection_pool_size
-        )
-        set_utility(Utility.MAINDB_DRIVER, tikv_driver)
-    elif settings.driver == DriverConfig.PG:
+    if settings.driver == DriverConfig.PG:
         if not PG:
-            raise ConfigurationError("`asyncpg` python package not installed.")
+            raise ConfigurationError("`psycopg` python package not installed.")
         if settings.driver_pg_url is None:
             raise ConfigurationError("No DRIVER_PG_URL env var defined.")
         pg_driver = PGDriver(
             url=settings.driver_pg_url,
+            connection_pool_min_size=settings.driver_pg_connection_pool_min_size,
             connection_pool_max_size=settings.driver_pg_connection_pool_max_size,
+            acquire_timeout_ms=settings.driver_pg_connection_pool_acquire_timeout_ms,
         )
         set_utility(Utility.MAINDB_DRIVER, pg_driver)
     elif settings.driver == DriverConfig.LOCAL:
@@ -97,13 +68,11 @@ async def setup_driver() -> Driver:
             raise ConfigurationError("`aiofiles` python package not installed.")
         if settings.driver_local_url is None:
             raise ConfigurationError("No DRIVER_LOCAL_URL env var defined.")
-
         local_driver = LocalDriver(settings.driver_local_url)
         set_utility(Utility.MAINDB_DRIVER, local_driver)
+        warn("LocalDriver is not recommended for production use", RuntimeWarning)
     else:
-        raise ConfigurationError(
-            f"Invalid DRIVER defined configured: {settings.driver}"
-        )
+        raise ConfigurationError(f"Invalid DRIVER defined configured: {settings.driver}")
 
     driver = get_driver()
     if not driver.initialized: