datajunction-server 0.0.1a48__tar.gz → 0.0.1a49.dev0__tar.gz

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

@@ -1,6 +1,6 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.1
 Name: datajunction-server
-Version: 0.0.1a48
+Version: 0.0.1a49.dev0
 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.dev0/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.dev0/datajunction_server/__about__.py

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

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

@@ -5,7 +5,7 @@ Data related APIs.
 from http import HTTPStatus
 from typing import Annotated, Dict, List, Optional
 
-from fastapi import Depends, Header, Query, Request
+from fastapi import BackgroundTasks, Depends, Header, Query, Request
 from fastapi.responses import JSONResponse
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import joinedload, selectinload
@@ -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
@@ -181,11 +181,12 @@ async def get_data(  # pylint: disable=too-many-locals
     validate_access: access.ValidateAccessFn = Depends(  # pylint: disable=W0621
         validate_access,
     ),
+    background_tasks: BackgroundTasks,
 ) -> QueryWithResults:
     """
     Gets data for a node
     """
-    query = await get_sql(
+    query, query_request = await get_node_sql(
         node_name,
         dimensions,
         filters,
@@ -196,6 +197,7 @@ async def get_data(  # pylint: disable=too-many-locals
         engine_version=engine_version,
         current_user=current_user,
         validate_access=validate_access,
+        background_tasks=background_tasks,
     )
     node = await Node.get_by_name(session, node_name, raise_if_not_exists=True)
     available_engines = node.current.catalog.engines  # type: ignore
@@ -221,12 +223,108 @@ 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,
+    ),
+    background_tasks: BackgroundTasks,
+) -> 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,
+        background_tasks=background_tasks,
+    )
+    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.1a48 → datajunction_server-0.0.1a49.dev0}/datajunction_server/api/helpers.py

@@ -21,6 +21,7 @@ from datajunction_server.construction.build import (
     build_materialized_cube_node,
     build_metric_nodes,
     build_node,
+    get_default_criteria,
     rename_columns,
     validate_shared_dimensions,
 )
@@ -50,10 +51,9 @@ from datajunction_server.errors import (
 from datajunction_server.internal.engines import get_engine
 from datajunction_server.models import access
 from datajunction_server.models.attribute import RESERVED_ATTRIBUTE_NAMESPACE
-from datajunction_server.models.engine import Dialect
 from datajunction_server.models.history import status_change_history
 from datajunction_server.models.metric import TranslatedSQL
-from datajunction_server.models.node import BuildCriteria, NodeStatus
+from datajunction_server.models.node import NodeStatus
 from datajunction_server.models.node_type import NodeType
 from datajunction_server.models.query import ColumnMetadata, QueryWithResults
 from datajunction_server.naming import LOOKUP_CHARS
@@ -204,19 +204,7 @@ async def get_query(  # pylint: disable=too-many-arguments
     Get a query for a metric, dimensions, and filters
     """
     node = await Node.get_by_name(session, node_name)
-
-    # Builds the node for the engine's dialect if one is set or defaults to Spark
-    if (
-        not engine
-        and node.current  # type: ignore
-        and node.current.catalog  # type: ignore
-        and node.current.catalog.engines  # type: ignore
-    ):
-        engine = node.current.catalog.engines[0]  # type: ignore
-    build_criteria = BuildCriteria(
-        dialect=(engine.dialect if engine and engine.dialect else Dialect.SPARK),
-    )
-
+    build_criteria = get_default_criteria(node.current, engine)  # type: ignore
     query_ast = await build_node(
         session=session,
         node=node.current,  # type: ignore

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

@@ -26,6 +26,7 @@ from datajunction_server.api.helpers import (
 from datajunction_server.api.namespaces import create_node_namespace
 from datajunction_server.api.tags import get_tags_by_name
 from datajunction_server.constants import NODE_LIST_MAX
+from datajunction_server.database import DimensionLink
 from datajunction_server.database.attributetype import ColumnAttribute
 from datajunction_server.database.column import Column
 from datajunction_server.database.history import ActivityType, EntityType, History
@@ -103,7 +104,7 @@ from datajunction_server.sql.dag import (
     get_downstream_nodes,
     get_upstream_nodes,
 )
-from datajunction_server.sql.parsing.backends.antlr4 import parse
+from datajunction_server.sql.parsing.backends.antlr4 import parse, parse_rule
 from datajunction_server.utils import (
     Version,
     get_current_user,
@@ -967,11 +968,10 @@ async def refresh_source_node(
     if new_columns:
         # check if any of the columns have changed (only continue with update if they have)
         column_changes = {col.identifier() for col in current_revision.columns} != {
-            col.identifier() for col in new_columns
+            (col.name, str(parse_rule(str(col.type), "dataType")))
+            for col in new_columns
         }
 
-        # FIXME: there is a bug with type translation (bigint != long) - fix it. # pylint: disable=fixme
-
         # if the columns haven't changed and the node has a table, we can skip the update
         if not column_changes:
             if not source_node.missing_table:  # type: ignore
@@ -998,6 +998,16 @@ async def refresh_source_node(
         schema_=current_revision.schema_,
         table=current_revision.table,
         status=current_revision.status,
+        dimension_links=[
+            DimensionLink(
+                dimension_id=link.dimension_id,
+                join_sql=link.join_sql,
+                join_type=link.join_type,
+                join_cardinality=link.join_cardinality,
+                materialization_conf=link.materialization_conf,
+            )
+            for link in current_revision.dimension_links
+        ],
     )
     new_revision.version = str(old_version.next_major_version())
     new_revision.columns = [

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

@@ -3,9 +3,9 @@
 SQL related APIs.
 """
 import logging
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
-from fastapi import Depends, Query
+from fastapi import BackgroundTasks, Depends, Query
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from datajunction_server.api.helpers import (
@@ -15,12 +15,14 @@ from datajunction_server.api.helpers import (
     validate_orderby,
 )
 from datajunction_server.construction.build import get_measures_query
+from datajunction_server.database import Engine
 from datajunction_server.database.queryrequest import QueryBuildType, QueryRequest
 from datajunction_server.database.user import User
 from datajunction_server.internal.access.authentication.http import SecureAPIRouter
 from datajunction_server.internal.access.authorization import validate_access
 from datajunction_server.internal.engines import get_engine
 from datajunction_server.models import access
+from datajunction_server.models.access import AccessControlStore
 from datajunction_server.models.metric import TranslatedSQL
 from datajunction_server.models.user import UserOutput
 from datajunction_server.utils import get_current_user, get_session, get_settings
@@ -108,12 +110,52 @@ 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 build_and_save_node_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: Engine,
+    access_control: AccessControlStore,
+) -> QueryRequest:
+    """
+    Build node SQL and save it to query requests
+    """
+    query_ast = await get_query(
+        session=session,
+        node_name=node_name,
+        dimensions=dimensions,
+        filters=filters,
+        orderby=orderby,
+        limit=limit,
+        engine=engine,
+        access_control=access_control,
+    )
+    columns = [
+        assemble_column_metadata(col)  # type: ignore
+        for col in query_ast.select.projection
+    ]
+    query = str(query_ast)
+    query_request = await QueryRequest.save_query_request(
+        session=session,
+        nodes=[node_name],
+        dimensions=dimensions,
+        filters=filters,
+        orderby=orderby,
+        limit=limit,
+        engine_name=engine.name if engine else None,
+        engine_version=engine.version if engine else None,
+        query_type=QueryBuildType.NODE,
+        query=query,
+        columns=[col.dict() for col in columns],
+    )
+    return query_request
+
+
+async def get_node_sql(  # pylint: disable=too-many-locals
     node_name: str,
     dimensions: List[str] = Query([]),
     filters: List[str] = Query([]),
@@ -123,11 +165,10 @@ 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
+    background_tasks: BackgroundTasks,
+) -> Tuple[TranslatedSQL, QueryRequest]:
     """
     Return SQL for a node.
     """
@@ -156,45 +197,85 @@ 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,
+        # Update the node SQL in a background task to keep it up-to-date
+        background_tasks.add_task(
+            build_and_save_node_sql,
+            node_name=node_name,
+            dimensions=dimensions,
+            filters=filters,
+            orderby=orderby,
+            limit=limit,
+            session=session,
+            engine=engine,
+            access_control=access_control,
+        )
+        return (
+            TranslatedSQL(
+                sql=query_request.query,
+                columns=query_request.columns,
+                dialect=engine.dialect if engine else None,
+            ),
+            query_request,
         )
 
-    query_ast = await get_query(
-        session=session,
+    query_request = await build_and_save_node_sql(
         node_name=node_name,
         dimensions=dimensions,
         filters=filters,
         orderby=orderby,
         limit=limit,
-        engine=engine,
+        session=session,
+        engine=engine,  # type: ignore
         access_control=access_control,
     )
-    columns = [
-        assemble_column_metadata(col)  # type: ignore
-        for col in query_ast.select.projection
-    ]
-    query = str(query_ast)
-    await QueryRequest.save_query_request(
-        session=session,
-        nodes=[node_name],
-        dimensions=dimensions,
-        filters=filters,
-        orderby=orderby,
-        limit=limit,
-        engine_name=engine.name if engine else None,
-        engine_version=engine.version if engine else None,
-        query_type=QueryBuildType.NODE,
-        query=query,
-        columns=[col.dict() for col in columns],
+    return (
+        TranslatedSQL(
+            sql=query_request.query,
+            columns=query_request.columns,
+            dialect=engine.dialect if engine else None,
+        ),
+        query_request,
     )
-    return TranslatedSQL(
-        sql=query,
-        columns=columns,
-        dialect=engine.dialect if engine else None,
+
+
+@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,
+    ),
+    background_tasks: BackgroundTasks,
+) -> 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,
+        background_tasks=background_tasks,
     )
+    return translated_sql
 
 
 @router.get("/sql/", response_model=TranslatedSQL, name="Get SQL For Metrics")