datajunction-server 0.0.1a95.dev0__tar.gz → 0.0.1a96__tar.gz

{datajunction_server-0.0.1a95.dev0 → datajunction_server-0.0.1a96}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: datajunction-server
-Version: 0.0.1a95.dev0
+Version: 0.0.1a96
 Summary: DataJunction server library for running to a DataJunction server
 Project-URL: Homepage, https://datajunction.io
 Project-URL: Repository, https://github.com/DataJunction/dj

datajunction_server-0.0.1a96/datajunction_server/__about__.py

@@ -0,0 +1,5 @@
+"""
+Version for Hatch
+"""
+
+__version__ = "0.0.1a96"

{datajunction_server-0.0.1a95.dev0 → datajunction_server-0.0.1a96}/datajunction_server/api/metrics.py

@@ -5,7 +5,7 @@ Metric related APIs.
 from http import HTTPStatus
 from typing import List, Optional
 
-from fastapi import Depends, HTTPException, Query
+from fastapi import BackgroundTasks, Depends, HTTPException, Query
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import selectinload
@@ -15,6 +15,8 @@ from datajunction_server.api.nodes import list_nodes
 from datajunction_server.database.node import Node, NodeRevision
 from datajunction_server.database.user import User
 from datajunction_server.errors import DJError, DJInvalidInputException, ErrorCode
+from datajunction_server.internal.caching.cachelib_cache import get_cache
+from datajunction_server.internal.caching.interface import Cache
 from datajunction_server.internal.access.authentication.http import SecureAPIRouter
 from datajunction_server.internal.access.authorization import validate_access
 from datajunction_server.models import access
@@ -69,17 +71,23 @@ async def list_metrics(
     validate_access: access.ValidateAccessFn = Depends(
         validate_access,
     ),
+    cache: Cache = Depends(get_cache),
+    background_tasks: BackgroundTasks,
 ) -> List[str]:
     """
     List all available metrics.
     """
-    return await list_nodes(
-        node_type=NodeType.METRIC,
-        prefix=prefix,
-        session=session,
-        current_user=current_user,
-        validate_access=validate_access,
-    )
+    metrics = cache.get("metrics")
+    if metrics is None:
+        metrics = await list_nodes(
+            node_type=NodeType.METRIC,
+            prefix=prefix,
+            session=session,
+            current_user=current_user,
+            validate_access=validate_access,
+        )
+        background_tasks.add_task(cache.set, "metrics", metrics)
+    return metrics
 
 
 @router.get("/metrics/metadata")

{datajunction_server-0.0.1a95.dev0 → datajunction_server-0.0.1a96}/datajunction_server/database/node.py

@@ -56,7 +56,7 @@ from datajunction_server.models.node_type import NodeType
 from datajunction_server.models.partition import PartitionType
 from datajunction_server.naming import amenable_name
 from datajunction_server.typing import UTCDatetime
-from datajunction_server.utils import SEPARATOR
+from datajunction_server.utils import SEPARATOR, execute_with_retry
 
 if TYPE_CHECKING:
     from datajunction_server.database.dimensionlink import DimensionLink
@@ -487,7 +487,7 @@ class Node(Base):
 
         limit = limit if limit and limit > 0 else 100
         statement = statement.limit(limit)
-        result = await session.execute(statement.options(*(options or [])))
+        result = await execute_with_retry(session, statement.options(*(options or [])))
         nodes = result.unique().scalars().all()
 
         # Reverse for backward pagination

{datajunction_server-0.0.1a95.dev0 → datajunction_server-0.0.1a96}/datajunction_server/errors.py

@@ -255,6 +255,15 @@ class DJNotImplementedException(DJException):
     http_status_code: int = 500
 
 
+class DJDatabaseException(DJException):
+    """
+    Exception raised when the database returns an error.
+    """
+
+    dbapi_exception: DBAPIExceptions = "DatabaseError"
+    http_status_code: int = 500
+
+
 class DJInternalErrorException(DJException):
     """
     Exception raised when we do something wrong in the code.

{datajunction_server-0.0.1a95.dev0 → datajunction_server-0.0.1a96}/datajunction_server/utils.py

@@ -16,7 +16,7 @@ from fastapi import Depends
 from rich.logging import RichHandler
 from sqlalchemy import AsyncAdaptedQueuePool
 from sqlalchemy.dialects.postgresql import insert
-from sqlalchemy.exc import MissingGreenlet
+from sqlalchemy.exc import MissingGreenlet, OperationalError
 from sqlalchemy.ext.asyncio import (
     AsyncEngine,
     AsyncSession,
@@ -24,6 +24,8 @@ from sqlalchemy.ext.asyncio import (
     async_sessionmaker,
     create_async_engine,
 )
+from sqlalchemy.sql import Select
+
 from starlette.requests import Request
 from yarl import URL
 
@@ -32,6 +34,7 @@ from datajunction_server.database.user import User
 from datajunction_server.enum import StrEnum
 from datajunction_server.errors import (
     DJAuthenticationException,
+    DJDatabaseException,
     DJInternalErrorException,
     DJInvalidInputException,
     DJUninitializedResourceException,
@@ -185,6 +188,29 @@ async def refresh_if_needed(session: AsyncSession, obj, attributes: list[str]):
         await session.refresh(obj, attributes_to_refresh)
 
 
+async def execute_with_retry(
+    session: AsyncSession,
+    statement: Select,
+    retries: int = 3,
+    base_delay: float = 1.0,
+):
+    """
+    Execute a SQLAlchemy statement with retry logic for transient errors.
+    """
+    attempt = 0
+    while True:
+        try:
+            return await session.execute(statement)
+        except OperationalError as exc:
+            attempt += 1
+            if attempt > retries:
+                raise DJDatabaseException(
+                    "Database operation failed after retries",
+                ) from exc
+            delay = base_delay * (2 ** (attempt - 1))
+            await asyncio.sleep(delay)
+
+
 def get_query_service_client(
     request: Request = None,
 ) -> Optional[QueryServiceClient]:

{datajunction_server-0.0.1a95.dev0 → datajunction_server-0.0.1a96}/tests/api/metrics_test.py

@@ -3,6 +3,7 @@ Tests for the metrics API.
 """
 
 import pytest
+from pytest_mock import MockerFixture
 from httpx import AsyncClient
 from sqlalchemy import select, text
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -1334,3 +1335,31 @@ async def test_create_invalid_metric(module__client_with_roads: AsyncClient):
     assert response.json()["message"] == (
         "Metric has an invalid query. The following are not allowed: GROUP BY, HAVING, ORDER BY"
     )
+
+
+@pytest.mark.asyncio
+async def test_read_metrics_when_cached(
+    module__client: AsyncClient,
+    mocker: MockerFixture,
+) -> None:
+    """
+    Test ``GET /metrics/`` with mocked cache returning a list of strings.
+    """
+    mock_list_nodes = mocker.patch(
+        "datajunction_server.api.metrics.list_nodes",
+        return_value=["metric1", "metric2", "metric3"],
+    )
+
+    # Should be a cache miss, triggering a cache in a background task
+    response1 = await module__client.get(
+        "/metrics/",
+        headers={"Cache-Control": "no-cache"},
+    )
+    assert response1.status_code == 200
+
+    # Should be a cache hit
+    response2 = await module__client.get("/metrics/")
+    assert response2.status_code == 200
+
+    # list_nodes should only be called once since the second used the cache
+    mock_list_nodes.assert_called_once()

{datajunction_server-0.0.1a95.dev0 → datajunction_server-0.0.1a96}/tests/utils_test.py

@@ -3,11 +3,12 @@ Tests for ``datajunction_server.utils``.
 """
 
 import logging
-from unittest.mock import patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from pytest_mock import MockerFixture
 from sqlalchemy import select
+from sqlalchemy.exc import OperationalError
 from sqlalchemy.ext.asyncio import AsyncSession
 from starlette.background import BackgroundTasks
 from testcontainers.postgres import PostgresContainer
@@ -15,9 +16,10 @@ from yarl import URL
 
 from datajunction_server.config import Settings
 from datajunction_server.database.user import OAuthProvider, User
-from datajunction_server.errors import DJException
+from datajunction_server.errors import DJDatabaseException, DJException
 from datajunction_server.utils import (
     Version,
+    execute_with_retry,
     get_and_update_current_user,
     get_engine,
     get_issue_url,
@@ -170,3 +172,46 @@ async def test_get_and_update_current_user(session: AsyncSession):
     assert found_user.name == "djuser"
     assert found_user.email == "userfoo@datajunction.io"
     assert found_user.oauth_provider == "basic"
+
+
+@pytest.mark.asyncio
+async def test_execute_with_retry_success_after_flaky_connection():
+    """
+    Test that execute_with_retry succeeds after a flaky connection.
+    """
+    session = AsyncMock(spec=AsyncSession)
+    statement = MagicMock()
+
+    # Simulate flaky DB: first 2 calls raise OperationalError, 3rd returns success
+    mock_result = MagicMock()
+    mock_result.unique.return_value.scalars.return_value.all.return_value = [
+        "node1",
+        "node2",
+    ]
+    session.execute.side_effect = [
+        OperationalError("flaky", None, None),
+        OperationalError("still flaky", None, None),
+        mock_result,
+    ]
+
+    result = await execute_with_retry(session, statement, retries=5, base_delay=0.01)
+    values = result.unique().scalars().all()
+    assert values == ["node1", "node2"]
+    assert session.execute.call_count == 3
+
+
+@pytest.mark.asyncio
+async def test_execute_with_retry_exhausts_retries():
+    """
+    Test that execute_with_retry exhausts retries and fails.
+    """
+    session = AsyncMock(spec=AsyncSession)
+    statement = MagicMock()
+
+    # Always fail
+    session.execute.side_effect = OperationalError("permanent fail", None, None)
+
+    with pytest.raises(DJDatabaseException):
+        await execute_with_retry(session, statement, retries=3, base_delay=0.01)
+
+    assert session.execute.call_count == 4  # initial try + 3 retries

datajunction_server-0.0.1a95.dev0/datajunction_server/__about__.py

@@ -1,5 +0,0 @@
-"""
-Version for Hatch
-"""
-
-__version__ = "0.0.1a95.dev0"