datajunction-server 0.0.1a47.dev0__tar.gz → 0.0.1a49__tar.gz

{datajunction_server-0.0.1a47.dev0 → datajunction_server-0.0.1a49}/Makefile

@@ -18,7 +18,7 @@ docker-run:
 	docker compose up
 
 test:
-	pdm run pytest -n auto --cov=datajunction_server --cov-report=html -vv tests/ --doctest-modules datajunction_server --without-integration --without-slow-integration ${PYTEST_ARGS}
+	pdm run pytest -n auto --cov=datajunction_server -vv tests/ --doctest-modules datajunction_server --without-integration --without-slow-integration ${PYTEST_ARGS}
 
 integration:
 	pdm run pytest --cov=dj -vv tests/ --doctest-modules datajunction_server --with-integration --with-slow-integration

{datajunction_server-0.0.1a47.dev0 → datajunction_server-0.0.1a49}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: datajunction-server
-Version: 0.0.1a47.dev0
+Version: 0.0.1a49
 Summary: DataJunction server library for running to a DataJunction server
 Project-URL: repository, https://github.com/DataJunction/dj
 Author-email: DataJunction Authors <yian.shang@gmail.com>

datajunction_server-0.0.1a49/alembic/versions/2024_05_21_0012-57fc93ef6947_add_query_id_to_queryrequest.py

@@ -0,0 +1,28 @@
+"""Add query_id to queryrequest
+
+Revision ID: 57fc93ef6947
+Revises: 9b1227ff17f4
+Create Date: 2024-05-21 00:12:11.303914+00:00
+
+"""
+# pylint: disable=no-member, invalid-name, missing-function-docstring, unused-import, no-name-in-module
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "57fc93ef6947"
+down_revision = "9b1227ff17f4"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    with op.batch_alter_table("queryrequest", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("query_id", sa.String(), nullable=True))
+
+
+def downgrade():
+    with op.batch_alter_table("queryrequest", schema=None) as batch_op:
+        batch_op.drop_column("query_id")

datajunction_server-0.0.1a49/datajunction_server/__about__.py

@@ -0,0 +1,4 @@
+"""
+Version for Hatch
+"""
+__version__ = "0.0.1a49"

{datajunction_server-0.0.1a47.dev0 → datajunction_server-0.0.1a49}/datajunction_server/api/data.py

@@ -15,7 +15,7 @@ from datajunction_server.api.helpers import (
     build_sql_for_multiple_metrics,
     query_event_stream,
 )
-from datajunction_server.api.sql import get_sql
+from datajunction_server.api.sql import get_node_sql
 from datajunction_server.database.availabilitystate import AvailabilityState
 from datajunction_server.database.history import ActivityType, EntityType, History
 from datajunction_server.database.node import Node, NodeRevision
@@ -185,7 +185,7 @@ async def get_data(  # pylint: disable=too-many-locals
     """
     Gets data for a node
     """
-    query = await get_sql(
+    query, query_request = await get_node_sql(
         node_name,
         dimensions,
         filters,
@@ -221,12 +221,106 @@ async def get_data(  # pylint: disable=too-many-locals
         query_create,
         headers={"Cache-Control": cache_control},
     )
+    query_request.query_id = result.id
+
     # Inject column info if there are results
     if result.results.__root__:  # pragma: no cover
-        result.results.__root__[0].columns = query.columns
+        result.results.__root__[0].columns = query.columns  # type: ignore
     return result
 
 
+@router.get("/stream/{node_name}", response_model=QueryWithResults)
+async def get_data_stream_for_node(  # pylint: disable=R0914, R0913
+    node_name: str,
+    *,
+    dimensions: List[str] = Query([], description="Dimensional attributes to group by"),
+    filters: List[str] = Query([], description="Filters on dimensional attributes"),
+    orderby: List[str] = Query([], description="Expression to order by"),
+    limit: Optional[int] = Query(
+        None,
+        description="Number of rows to limit the data retrieved to",
+    ),
+    cache_control: Annotated[str, Header()] = "",
+    session: AsyncSession = Depends(get_session),
+    request: Request,
+    query_service_client: QueryServiceClient = Depends(get_query_service_client),
+    engine_name: Optional[str] = None,
+    engine_version: Optional[str] = None,
+    current_user: Optional[User] = Depends(get_current_user),
+    validate_access: access.ValidateAccessFn = Depends(  # pylint: disable=W0621
+        validate_access,
+    ),
+) -> QueryWithResults:
+    """
+    Return data for a node using server side events
+    """
+    query, query_request = await get_node_sql(
+        node_name,
+        dimensions,
+        filters,
+        orderby,
+        limit,
+        session=session,
+        engine_name=engine_name,
+        engine_version=engine_version,
+        current_user=current_user,
+        validate_access=validate_access,
+    )
+    if query_request and query_request.query_id:
+        return EventSourceResponse(
+            query_event_stream(
+                query=QueryWithResults(
+                    id=query_request.query_id,
+                    submitted_query=query_request.query,
+                    results=[],
+                    errors=[],
+                ),
+                query_service_client=query_service_client,
+                columns=query.columns,  # type: ignore
+                request=request,
+            ),
+        )
+
+    node = await Node.get_by_name(session, node_name, raise_if_not_exists=True)
+    available_engines = node.current.catalog.engines  # type: ignore
+    engine = (
+        await get_engine(session, engine_name, engine_version)  # type: ignore
+        if engine_name
+        else available_engines[0]
+    )
+    if engine not in available_engines:
+        raise DJInvalidInputException(  # pragma: no cover
+            f"The selected engine is not available for the node {node_name}. "
+            f"Available engines include: {', '.join(engine.name for engine in available_engines)}",
+        )
+
+    query_create = QueryCreate(
+        engine_name=engine.name,
+        catalog_name=node.current.catalog.name,  # type: ignore
+        engine_version=engine.version,
+        submitted_query=query.sql,
+        async_=True,
+    )
+    initial_query_info = query_service_client.submit_query(
+        query_create,
+        headers={"Cache-Control": cache_control or "max-age=86400"},
+    )
+
+    # Save the external query id reference
+    query_request.query_id = initial_query_info.id
+    session.add(query_request)
+    await session.commit()
+
+    return EventSourceResponse(
+        query_event_stream(
+            query=initial_query_info,
+            query_service_client=query_service_client,
+            columns=query.columns,  # type: ignore
+            request=request,
+        ),
+    )
+
+
 @router.get(
     "/data/query/{query_id}",
     response_model=QueryWithResults,

{datajunction_server-0.0.1a47.dev0 → datajunction_server-0.0.1a49}/datajunction_server/api/nodes.py

@@ -622,7 +622,6 @@ async def create_cube(
     """
     Create a cube node.
     """
-    print("data", data)
     await raise_if_node_exists(session, data.name)
 
     # if the node previously existed and now is inactive

{datajunction_server-0.0.1a47.dev0 → datajunction_server-0.0.1a49}/datajunction_server/api/sql.py

@@ -3,7 +3,7 @@
 SQL related APIs.
 """
 import logging
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
 from fastapi import Depends, Query
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -108,12 +108,7 @@ async def get_measures_sql_for_cube(
     return measures_query
 
 
-@router.get(
-    "/sql/{node_name}/",
-    response_model=TranslatedSQL,
-    name="Get SQL For A Node",
-)
-async def get_sql(  # pylint: disable=too-many-locals
+async def get_node_sql(  # pylint: disable=too-many-locals
     node_name: str,
     dimensions: List[str] = Query([]),
     filters: List[str] = Query([]),
@@ -123,11 +118,9 @@ async def get_sql(  # pylint: disable=too-many-locals
     session: AsyncSession = Depends(get_session),
     engine_name: Optional[str] = None,
     engine_version: Optional[str] = None,
-    current_user: Optional[User] = Depends(get_current_user),
-    validate_access: access.ValidateAccessFn = Depends(  # pylint: disable=W0621
-        validate_access,
-    ),
-) -> TranslatedSQL:
+    current_user: Optional[User],
+    validate_access: access.ValidateAccessFn,  # pylint: disable=redefined-outer-name
+) -> Tuple[TranslatedSQL, QueryRequest]:
     """
     Return SQL for a node.
     """
@@ -156,10 +149,13 @@ async def get_sql(  # pylint: disable=too-many-locals
         engine_version=engine.version if engine else None,
         query_type=QueryBuildType.NODE,
     ):
-        return TranslatedSQL(
-            sql=query_request.query,
-            columns=query_request.columns,
-            dialect=engine.dialect if engine else None,
+        return (
+            TranslatedSQL(
+                sql=query_request.query,
+                columns=query_request.columns,
+                dialect=engine.dialect if engine else None,
+            ),
+            query_request,
         )
 
     query_ast = await get_query(
@@ -177,7 +173,7 @@ async def get_sql(  # pylint: disable=too-many-locals
         for col in query_ast.select.projection
     ]
     query = str(query_ast)
-    await QueryRequest.save_query_request(
+    query_request = await QueryRequest.save_query_request(
         session=session,
         nodes=[node_name],
         dimensions=dimensions,
@@ -190,11 +186,52 @@ async def get_sql(  # pylint: disable=too-many-locals
         query=query,
         columns=[col.dict() for col in columns],
     )
-    return TranslatedSQL(
-        sql=query,
-        columns=columns,
-        dialect=engine.dialect if engine else None,
+    return (
+        TranslatedSQL(
+            sql=query,
+            columns=columns,
+            dialect=engine.dialect if engine else None,
+        ),
+        query_request,
+    )
+
+
+@router.get(
+    "/sql/{node_name}/",
+    response_model=TranslatedSQL,
+    name="Get SQL For A Node",
+)
+async def get_sql(  # pylint: disable=too-many-locals
+    node_name: str,
+    dimensions: List[str] = Query([]),
+    filters: List[str] = Query([]),
+    orderby: List[str] = Query([]),
+    limit: Optional[int] = None,
+    *,
+    session: AsyncSession = Depends(get_session),
+    engine_name: Optional[str] = None,
+    engine_version: Optional[str] = None,
+    current_user: Optional[User] = Depends(get_current_user),
+    validate_access: access.ValidateAccessFn = Depends(  # pylint: disable=W0621
+        validate_access,
+    ),
+) -> TranslatedSQL:
+    """
+    Return SQL for a node.
+    """
+    translated_sql, _ = await get_node_sql(
+        node_name,
+        dimensions,
+        filters,
+        orderby,
+        limit,
+        session=session,
+        engine_name=engine_name,
+        engine_version=engine_version,
+        current_user=current_user,
+        validate_access=validate_access,
     )
+    return translated_sql
 
 
 @router.get("/sql/", response_model=TranslatedSQL, name="Get SQL For Metrics")

{datajunction_server-0.0.1a47.dev0 → datajunction_server-0.0.1a49}/datajunction_server/construction/build.py

@@ -773,6 +773,9 @@ def _get_node_table(
     If a node has a materialization available, return the materialized table
     """
     table = None
+    can_use_materialization = (
+        build_criteria and node.name != build_criteria.target_node_name
+    )
     if node.type == NodeType.SOURCE:
         if node.table:
             name = ast.Name(
@@ -796,8 +799,12 @@ def _get_node_table(
             ],
             _dj_node=node,
         )
-    elif node.availability and node.availability.is_available(
-        criteria=build_criteria,
+    elif (
+        can_use_materialization
+        and node.availability
+        and node.availability.is_available(
+            criteria=build_criteria,
+        )
     ):  # pragma: no cover
         table = ast.Table(
             ast.Name(
@@ -822,6 +829,23 @@ def _get_node_table(
     return table
 
 
+def get_default_criteria(
+    node: NodeRevision,
+) -> BuildCriteria:
+    """
+    Get the default build criteria for a node.
+    """
+    return BuildCriteria(
+        # set the dialect by finding available engines for this node, or default to Spark
+        dialect=(
+            node.catalog.engines[0].dialect
+            if node.catalog and node.catalog.engines and node.catalog.engines[0].dialect
+            else Dialect.SPARK
+        ),
+        target_node_name=node.name,
+    )
+
+
 async def build_node(  # pylint: disable=too-many-arguments
     session: AsyncSession,
     node: NodeRevision,
@@ -844,15 +868,7 @@ async def build_node(  # pylint: disable=too-many-arguments
 
     # Set the dialect by finding available engines for this node, or default to Spark
     if not build_criteria:
-        build_criteria = BuildCriteria(
-            dialect=(
-                node.catalog.engines[0].dialect
-                if node.catalog
-                and node.catalog.engines
-                and node.catalog.engines[0].dialect
-                else Dialect.SPARK
-            ),
-        )
+        build_criteria = get_default_criteria(node)
 
     # get dimension columns which are required
     # in the stated bound dimensions on the metric node

{datajunction_server-0.0.1a47.dev0 → datajunction_server-0.0.1a49}/datajunction_server/database/availabilitystate.py

@@ -70,7 +70,8 @@ class AvailabilityState(Base):  # pylint: disable=too-few-public-methods
         """
         Determine whether an availability state is useable given criteria
         """
-        # Criteria to determine if an availability state should be used needs to be added
+        # TODO: we should evaluate this availability state against the criteria. # pylint: disable=fixme
+        #       Remember that VTTS can be also evaluated at runtime dependency.
         return True
 
 

{datajunction_server-0.0.1a47.dev0 → datajunction_server-0.0.1a49}/datajunction_server/database/queryrequest.py

@@ -159,6 +159,8 @@ class QueryRequest(Base):  # type: ignore  # pylint: disable=too-few-public-meth
         DateTime(timezone=True),
         default=partial(datetime.now, timezone.utc),
     )
+    # External identifier for the query
+    query_id: Mapped[Optional[str]]
 
     @classmethod
     async def get_query_request(
@@ -220,7 +222,7 @@ class QueryRequest(Base):  # type: ignore  # pylint: disable=too-few-public-meth
         query: str,
         columns: List[Dict[str, Any]],
         other_args: Optional[Dict[str, Any]] = None,
-    ):
+    ) -> "QueryRequest":
         """
         Retrieves saved query for a node SQL request
         """
@@ -266,6 +268,7 @@ class QueryRequest(Base):  # type: ignore  # pylint: disable=too-few-public-meth
             )
             session.add(query_request)
             await session.commit()
+        return query_request
 
     @classmethod
     async def to_versioned_query_request(  # pylint: disable=too-many-locals

{datajunction_server-0.0.1a47.dev0 → datajunction_server-0.0.1a49}/datajunction_server/internal/materializations.py

@@ -7,7 +7,11 @@ from sqlalchemy.exc import InvalidRequestError
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from datajunction_server.api.helpers import build_sql_for_multiple_metrics
-from datajunction_server.construction.build import build_node, get_measures_query
+from datajunction_server.construction.build import (
+    build_node,
+    get_default_criteria,
+    get_measures_query,
+)
 from datajunction_server.database.materialization import Materialization
 from datajunction_server.database.node import NodeRevision
 from datajunction_server.errors import DJException, DJInvalidInputException
@@ -183,11 +187,15 @@ async def build_non_cube_materialization_config(
     """
     Build materialization config for non-cube nodes (transforms and dimensions).
     """
+    build_criteria = get_default_criteria(
+        node=current_revision,
+    )
     materialization_ast = await build_node(
         session=session,
         node=current_revision,
         dimensions=[],
         orderby=[],
+        build_criteria=build_criteria,
     )
     generic_config = GenericMaterializationConfig(
         lookback_window=upsert.config.lookback_window,

{datajunction_server-0.0.1a47.dev0 → datajunction_server-0.0.1a49}/datajunction_server/models/node.py

@@ -45,7 +45,7 @@ class BuildCriteria:
 
     timestamp: Optional[UTCDatetime] = None
     dialect: Dialect = Dialect.SPARK
-    for_materialization: bool = False
+    target_node_name: Optional[str] = None
 
 
 class NodeMode(StrEnum):

{datajunction_server-0.0.1a47.dev0 → datajunction_server-0.0.1a49}/datajunction_server/sql/functions.py

@@ -278,6 +278,8 @@ class Aggregate(Function):
             ast,
         )
 
+        # aggregate's lambda function can take three or four arguments, depending on whether
+        # an optional finish function is provided
         if len(args) == 4:
             expr, start, merge, finish = args
         else:

{datajunction_server-0.0.1a47.dev0 → datajunction_server-0.0.1a49}/datajunction_server/sql/parsing/ast.py

@@ -1734,7 +1734,10 @@ class Function(Named, Operation):
         try:
             self.function().compile_lambda(*self.args)
         except DJParseException as parse_exc:
-            logger.warning(parse_exc)
+            if "Cannot resolve type of column" in parse_exc.message:
+                logger.warning(parse_exc)
+            else:
+                raise parse_exc
 
         for child in self.children:
             if not child.is_compiled():

{datajunction_server-0.0.1a47.dev0 → datajunction_server-0.0.1a49}/tests/api/data_test.py

@@ -13,7 +13,9 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import joinedload
 
 from datajunction_server.api.main import app
+from datajunction_server.database import QueryRequest
 from datajunction_server.database.node import Node, NodeRevision
+from datajunction_server.database.queryrequest import QueryBuildType
 from datajunction_server.internal.access.authorization import validate_access
 from datajunction_server.models import access
 from datajunction_server.models.node import AvailabilityStateBase
@@ -552,10 +554,13 @@ class TestDataForNode:
     @pytest.mark.asyncio
     async def test_stream_multiple_metrics_and_dimensions_data(
         self,
+        session: AsyncSession,
         client_with_query_service_example_loader,
     ) -> None:
         """
-        Test streaming query status for multiple metrics and dimensions
+        Test streaming query status for
+        (a) multiple metrics and dimensions and
+        (b) node data
         """
         custom_client = await client_with_query_service_example_loader(["ROADS"])
         async with custom_client.stream(
@@ -567,6 +572,64 @@ class TestDataForNode:
             },
         ) as response:
             assert response.status_code == 200
+            full_text = "".join([text async for text in response.aiter_text()])
+            assert "event: message" in full_text
+            assert "avg(repair_order_details.price)" in full_text
+
+        # Test streaming of node data for a metric
+        async with custom_client.stream(
+            "GET",
+            "/stream/default.num_repair_orders?dimensions=default.dispatcher.company_name&limit=10",
+            headers={
+                "Accept": "text/event-stream",
+            },
+        ) as response:
+            assert response.status_code == 200
+            full_text = "".join([text async for text in response.aiter_text()])
+            assert "event: message" in full_text
+            assert "count(default_DOT_repair_orders_fact.repair_order_id)" in full_text
+
+        # Test streaming of node data for a transform
+        async with custom_client.stream(
+            "GET",
+            "/stream/default.repair_orders_fact?"
+            "dimensions=default.dispatcher.company_name&limit=10",
+            headers={
+                "Accept": "text/event-stream",
+            },
+        ) as response:
+            assert response.status_code == 200
+            full_text = "\n".join([text async for text in response.aiter_lines()])
+            assert "event: message" in full_text
+            assert "SELECT  default_DOT_repair_orders_fact.repair_order_id" in full_text
+
+        # Check that the query request for the above transform has an external query id saved
+        query_request = await QueryRequest.get_query_request(
+            session=session,
+            query_type=QueryBuildType.NODE,
+            nodes=["default.repair_orders_fact"],
+            dimensions=["default.dispatcher.company_name"],
+            engine_name=None,
+            engine_version=None,
+            filters=[],
+            limit=10,
+            orderby=[],
+        )
+        assert query_request.query_id == "bd98d6be-e2d2-413e-94c7-96d9411ddee2"  # type: ignore
+
+        # Hit the same SSE stream again
+        async with custom_client.stream(
+            "GET",
+            "/stream/default.repair_orders_fact?"
+            "dimensions=default.dispatcher.company_name&limit=10",
+            headers={
+                "Accept": "text/event-stream",
+            },
+        ) as response:
+            assert response.status_code == 200
+            full_text = "\n".join([text async for text in response.aiter_lines()])
+            assert "event: message" in full_text
+            assert "SELECT  default_DOT_repair_orders_fact.repair_order_id" in full_text
 
     @pytest.mark.asyncio
     async def test_get_data_for_query_id(