datajunction-server 0.0.1a45.dev14__tar.gz → 0.0.1a46__tar.gz

{datajunction_server-0.0.1a45.dev14 → datajunction_server-0.0.1a46}/.pre-commit-config.yaml

@@ -90,7 +90,7 @@ repos:
   hooks:
   - id: pylint
     name: pylint
-    entry: pylint --disable=duplicate-code,use-implicit-booleaness-not-comparison,wrong-import-order,too-many-arguments
+    entry: pylint --disable=duplicate-code,use-implicit-booleaness-not-comparison,wrong-import-order,too-many-arguments,cyclic-import
     language: system
     types: [python]
     exclude: ^templates/

{datajunction_server-0.0.1a45.dev14 → datajunction_server-0.0.1a46}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: datajunction-server
-Version: 0.0.1a45.dev14
+Version: 0.0.1a46
 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.1a46/alembic/versions/2024_04_30_1556-de7ec1c82fe0_add_query_requests.py

@@ -0,0 +1,98 @@
+"""Add query requests
+
+Revision ID: de7ec1c82fe0
+Revises: d61fb7e48cc3
+Create Date: 2024-04-30 15:56:20.193978+00:00
+
+"""
+# pylint: disable=no-member, invalid-name, missing-function-docstring, unused-import, no-name-in-module
+
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "de7ec1c82fe0"
+down_revision = "d61fb7e48cc3"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.create_table(
+        "queryrequest",
+        sa.Column("id", sa.BigInteger(), nullable=False),
+        sa.Column(
+            "query_type",
+            sa.Enum("METRICS", "MEASURES", "NODE", name="querybuildtype"),
+            nullable=False,
+        ),
+        sa.Column(
+            "nodes",
+            postgresql.JSONB(astext_type=sa.Text()),
+            server_default=sa.text("'[]'::jsonb"),
+            nullable=False,
+        ),
+        sa.Column(
+            "parents",
+            postgresql.JSONB(astext_type=sa.Text()),
+            server_default=sa.text("'[]'::jsonb"),
+            nullable=False,
+        ),
+        sa.Column(
+            "dimensions",
+            postgresql.JSONB(astext_type=sa.Text()),
+            server_default=sa.text("'[]'::jsonb"),
+            nullable=False,
+        ),
+        sa.Column(
+            "filters",
+            postgresql.JSONB(astext_type=sa.Text()),
+            server_default=sa.text("'[]'::jsonb"),
+            nullable=False,
+        ),
+        sa.Column("limit", sa.Integer(), nullable=True),
+        sa.Column(
+            "orderby",
+            postgresql.JSONB(astext_type=sa.Text()),
+            server_default=sa.text("'[]'::jsonb"),
+            nullable=False,
+        ),
+        sa.Column("engine_name", sa.String(), nullable=True),
+        sa.Column("engine_version", sa.String(), nullable=True),
+        sa.Column(
+            "other_args",
+            postgresql.JSONB(astext_type=sa.Text()),
+            server_default=sa.text("'{}'::jsonb"),
+            nullable=False,
+        ),
+        sa.Column("query", sa.String(), nullable=False),
+        sa.Column(
+            "columns",
+            postgresql.JSONB(astext_type=sa.Text()),
+            server_default=sa.text("'[]'::jsonb"),
+            nullable=False,
+        ),
+        sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint(
+            "query_type",
+            "nodes",
+            "parents",
+            "dimensions",
+            "filters",
+            "engine_name",
+            "engine_version",
+            "limit",
+            "orderby",
+            name="query_request_unique",
+            postgresql_nulls_not_distinct=True,
+        ),
+    )
+
+
+def downgrade():
+    op.drop_table("queryrequest")
+    op.execute("DROP TYPE querybuildtype")

datajunction_server-0.0.1a46/datajunction_server/__about__.py

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

{datajunction_server-0.0.1a45.dev14 → datajunction_server-0.0.1a46}/datajunction_server/api/helpers.py

@@ -522,6 +522,10 @@ async def resolve_downstream_references(
                 .scalar_one()
             )
             await session.refresh(node_revision, ["node"])
+            await session.refresh(
+                downstream_node_revision,
+                ["parents", "missing_parents"],
+            )
             downstream_node_revision.parents.append(node_revision.node)
             downstream_node_revision.missing_parents.remove(missing_parent)
             node_validator = await validate_node_data(

{datajunction_server-0.0.1a45.dev14 → datajunction_server-0.0.1a46}/datajunction_server/api/sql.py

@@ -15,6 +15,7 @@ from datajunction_server.api.helpers import (
     validate_orderby,
 )
 from datajunction_server.construction.build import get_measures_query
+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
@@ -35,6 +36,13 @@ async def get_measures_sql_for_cube(
     dimensions: List[str] = Query([]),
     filters: List[str] = Query([]),
     *,
+    include_all_columns: bool = Query(
+        False,
+        description=(
+            "Whether to include all columns or only those necessary "
+            "for the metrics and dimensions in the cube"
+        ),
+    ),
     session: AsyncSession = Depends(get_session),
     engine_name: Optional[str] = None,
     engine_version: Optional[str] = None,
@@ -48,7 +56,30 @@ async def get_measures_sql_for_cube(
     This SQL can be used to produce an intermediate table with all the measures
     and dimensions needed for an analytics database (e.g., Druid).
     """
-    return await get_measures_query(
+    if query_request := await QueryRequest.get_query_request(
+        session,
+        nodes=metrics,
+        dimensions=dimensions,
+        filters=filters,
+        orderby=[],
+        limit=None,
+        engine_name=engine_name,
+        engine_version=engine_version,
+        query_type=QueryBuildType.MEASURES,
+        other_args={"include_all_columns": include_all_columns},
+    ):
+        engine = (
+            await get_engine(session, engine_name, engine_version)  # type: ignore
+            if engine_name
+            else None
+        )
+        return TranslatedSQL(
+            sql=query_request.query,
+            columns=query_request.columns,
+            dialect=engine.dialect if engine else None,
+        )
+
+    measures_query = await get_measures_query(
         session=session,
         metrics=metrics,
         dimensions=dimensions,
@@ -57,7 +88,24 @@ async def get_measures_sql_for_cube(
         engine_version=engine_version,
         current_user=current_user,
         validate_access=validate_access,
+        include_all_columns=include_all_columns,
+    )
+
+    await QueryRequest.save_query_request(
+        session=session,
+        nodes=metrics,
+        dimensions=dimensions,
+        filters=filters,
+        orderby=[],
+        limit=None,
+        engine_name=engine_name,
+        engine_version=engine_version,
+        query_type=QueryBuildType.MEASURES,
+        query=measures_query.sql,
+        columns=[col.dict() for col in measures_query.columns],  # type: ignore
+        other_args={"include_all_columns": include_all_columns},
     )
+    return measures_query
 
 
 @router.get(
@@ -65,7 +113,7 @@ async def get_measures_sql_for_cube(
     response_model=TranslatedSQL,
     name="Get SQL For A Node",
 )
-async def get_sql(
+async def get_sql(  # pylint: disable=too-many-locals
     node_name: str,
     dimensions: List[str] = Query([]),
     filters: List[str] = Query([]),
@@ -95,6 +143,24 @@ async def get_sql(
         else None
     )
     validate_orderby(orderby, [node_name], dimensions)
+
+    if query_request := await QueryRequest.get_query_request(
+        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,
+    ):
+        return TranslatedSQL(
+            sql=query_request.query,
+            columns=query_request.columns,
+            dialect=engine.dialect if engine else None,
+        )
+
     query_ast = await get_query(
         session=session,
         node_name=node_name,
@@ -109,8 +175,22 @@ async def get_sql(
         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=str(query_ast),
+        sql=query,
         columns=columns,
         dialect=engine.dialect if engine else None,
     )
@@ -142,6 +222,28 @@ async def get_sql_for_metrics(
         base_verb=access.ResourceRequestVerb.READ,
     )
 
+    if query_request := await QueryRequest.get_query_request(
+        session,
+        nodes=metrics,
+        dimensions=dimensions,
+        filters=filters,
+        orderby=orderby,
+        limit=limit,
+        engine_name=engine_name,
+        engine_version=engine_version,
+        query_type=QueryBuildType.METRICS,
+    ):
+        engine = (
+            await get_engine(session, engine_name, engine_version)  # type: ignore
+            if engine_name
+            else None
+        )
+        return TranslatedSQL(
+            sql=query_request.query,
+            columns=query_request.columns,
+            dialect=engine.dialect if engine else None,
+        )
+
     translated_sql, _, _ = await build_sql_for_multiple_metrics(
         session,
         metrics,
@@ -153,4 +255,18 @@ async def get_sql_for_metrics(
         engine_version,
         access_control,
     )
+
+    await QueryRequest.save_query_request(
+        session=session,
+        nodes=metrics,
+        dimensions=dimensions,
+        filters=filters,
+        orderby=orderby,
+        limit=limit,
+        engine_name=engine_name,
+        engine_version=engine_version,
+        query_type=QueryBuildType.METRICS,
+        query=translated_sql.sql,
+        columns=[col.dict() for col in translated_sql.columns],  # type: ignore
+    )
     return translated_sql

{datajunction_server-0.0.1a45.dev14 → datajunction_server-0.0.1a46}/datajunction_server/construction/build.py

@@ -426,31 +426,42 @@ async def _build_tables_on_select(
                     filter_asts = (  # pylint: disable=consider-using-ternary
                         node_query.select.where and [node_query.select.where] or []
                     )
+                    foreign_keys_map = {
+                        left.alias_or_name.name: right
+                        for link in node.dimension_links
+                        for left, right in link.foreign_key_mapping().items()
+                    }
+                    foreign_keys_alias_map = {
+                        col.alias_or_name.name: col  # type: ignore
+                        for col in node_query.select.projection
+                    }
                     for filter_ in filters:
                         temp_select = parse(f"select * where {filter_}").select
                         referenced_cols = list(temp_select.find_all(ast.Column))
 
                         # We can only push down the filter if all columns referenced by the filter
                         # are available as foreign key columns on the node
-                        foreign_keys_map = {
-                            left.alias_or_name.name: right
-                            for link in node.dimension_links
-                            for left, right in link.foreign_key_mapping().items()
-                        }
                         if all(
                             col.alias_or_name.name in fk_column_mapping
                             or col.alias_or_name.name in foreign_keys_map
                             for col in referenced_cols
                         ):
+                            # Renames the columns to the foreign key columns
                             for col in referenced_cols:
                                 ref_col_name = (
                                     fk_column_mapping[col.alias_or_name.name].name
                                     if col.alias_or_name.name in fk_column_mapping
                                     else foreign_keys_map[
                                         col.alias_or_name.name
-                                    ].alias_or_name.name
+                                    ].alias_or_name
                                 )
                                 col.name = ast.Name(name=ref_col_name)
+                                if (  # pragma: no cover
+                                    col.alias_or_name.name in foreign_keys_alias_map
+                                ):
+                                    col.name = foreign_keys_alias_map[  # type: ignore
+                                        col.alias_or_name.name
+                                    ].name
                             filter_asts.append(
                                 temp_select.where,  # type:ignore
                             )
@@ -1469,6 +1480,7 @@ async def get_measures_query(
     current_user: Optional[User] = None,
     validate_access: access.ValidateAccessFn = None,
     cast_timestamp_to_ms: bool = False,
+    include_all_columns: bool = False,
 ) -> TranslatedSQL:
     """
     Builds the measures SQL for a set of metrics with dimensions and filters.
@@ -1535,18 +1547,19 @@ async def get_measures_query(
         )
 
         # Select only columns that were one of the necessary measures
-        parent_ast.select.projection = [
-            expr
-            for expr in parent_ast.select.projection
-            if from_amenable_name(expr.alias_or_name.identifier(False)).split(  # type: ignore
-                SEPARATOR,
-            )[
-                -1
+        if not include_all_columns:
+            parent_ast.select.projection = [
+                expr
+                for expr in parent_ast.select.projection
+                if from_amenable_name(expr.alias_or_name.identifier(False)).split(  # type: ignore
+                    SEPARATOR,
+                )[
+                    -1
+                ]
+                in parents_to_measures[parent_node.name]
+                or from_amenable_name(expr.alias_or_name.identifier(False))  # type: ignore
+                in dimensions_without_roles
             ]
-            in parents_to_measures[parent_node.name]
-            or from_amenable_name(expr.alias_or_name.identifier(False))  # type: ignore
-            in dimensions_without_roles
-        ]
         await session.refresh(parent_node.current, ["columns"])
         parent_ast = rename_columns(parent_ast, parent_node.current)
 

{datajunction_server-0.0.1a45.dev14 → datajunction_server-0.0.1a46}/datajunction_server/database/__init__.py

@@ -12,6 +12,7 @@ __all__ = [
     "NodeNamespace",
     "NodeRevision",
     "Partition",
+    "QueryRequest",
     "Table",
     "Tag",
     "User",
@@ -28,6 +29,7 @@ from datajunction_server.database.measure import Measure
 from datajunction_server.database.namespace import NodeNamespace
 from datajunction_server.database.node import Node, NodeRevision
 from datajunction_server.database.partition import Partition
+from datajunction_server.database.queryrequest import QueryRequest
 from datajunction_server.database.tag import Tag
 from datajunction_server.database.user import User
 from datajunction_server.models.history import History