datajunction-server 0.0.1a75__tar.gz → 0.0.1a76.dev0__tar.gz

{datajunction_server-0.0.1a75 → datajunction_server-0.0.1a76.dev0}/PKG-INFO

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

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

{datajunction_server-0.0.1a75 → datajunction_server-0.0.1a76.dev0}/datajunction_server/api/helpers.py

@@ -59,7 +59,7 @@ from datajunction_server.naming import LOOKUP_CHARS
 from datajunction_server.service_clients import QueryServiceClient
 from datajunction_server.sql.parsing import ast
 from datajunction_server.typing import END_JOB_STATES
-from datajunction_server.utils import SEPARATOR
+from datajunction_server.utils import SEPARATOR, refresh_if_needed
 
 _logger = logging.getLogger(__name__)
 
@@ -737,7 +737,7 @@ async def build_sql_for_multiple_metrics(  # pylint: disable=too-many-arguments,
     ]
     upstream_tables = [tbl for tbl in query_ast.find_all(ast.Table) if tbl.dj_node]
     for tbl in upstream_tables:
-        await session.refresh(tbl.dj_node, ["availability"])
+        await refresh_if_needed(session, tbl.dj_node, ["availability"])
     return (
         TranslatedSQL(
             sql=str(query_ast),

{datajunction_server-0.0.1a75 → datajunction_server-0.0.1a76.dev0}/datajunction_server/construction/build_v2.py

@@ -29,8 +29,8 @@ from datajunction_server.models.node_type import NodeType
 from datajunction_server.models.sql import GeneratedSQL
 from datajunction_server.naming import amenable_name, from_amenable_name
 from datajunction_server.sql.parsing.ast import CompileContext
-from datajunction_server.sql.parsing.backends.antlr4 import ast, parse
-from datajunction_server.utils import SEPARATOR
+from datajunction_server.sql.parsing.backends.antlr4 import ast, cached_parse, parse
+from datajunction_server.utils import SEPARATOR, refresh_if_needed
 
 logger = logging.getLogger(__name__)
 
@@ -189,7 +189,7 @@ async def get_measures_query(  # pylint: disable=too-many-locals
                 or from_amenable_name(expr.alias_or_name.identifier(False))  # type: ignore
                 in dimensions_without_roles
             ]
-        await session.refresh(parent_node.current, ["columns"])
+        await refresh_if_needed(session, parent_node.current, ["columns"])
         parent_ast = rename_columns(parent_ast, parent_node.current)
 
         # Sort the selected columns into dimension vs measure columns and
@@ -280,7 +280,11 @@ class QueryBuilder:  # pylint: disable=too-many-instance-attributes,too-many-pub
         """
         Create a QueryBuilder instance for the node revision.
         """
-        await session.refresh(node_revision, ["required_dimensions", "dimension_links"])
+        await refresh_if_needed(
+            session,
+            node_revision,
+            ["required_dimensions", "dimension_links"],
+        )
         instance = cls(session, node_revision, use_materialized=use_materialized)
         return instance
 
@@ -429,7 +433,11 @@ class QueryBuilder:  # pylint: disable=too-many-instance-attributes,too-many-pub
         7. Add all requested dimensions to the final select.
         8. Add order by and limit to the final select (TODO)
         """
-        await self.session.refresh(self.node_revision, ["availability", "columns"])
+        await refresh_if_needed(
+            self.session,
+            self.node_revision,
+            ["availability", "columns"],
+        )
         node_ast = (
             await compile_node_ast(self.session, self.node_revision)
             if not self.physical_table
@@ -463,7 +471,7 @@ class QueryBuilder:  # pylint: disable=too-many-instance-attributes,too-many-pub
         """
         Build the ORDER BY clause from the provided order expressions
         """
-        temp_orderbys = parse(
+        temp_orderbys = cached_parse(
             f"SELECT 1 ORDER BY {','.join(self._orderby)}",
         ).select.organization.order
         valid_sort_items = [
@@ -784,8 +792,8 @@ class CubeQueryBuilder:  # pylint: disable=too-many-instance-attributes
         Create a QueryBuilder instance for the node revision.
         """
         for node in metric_nodes:
-            await session.refresh(node, ["current"])
-            await session.refresh(node.current, ["required_dimensions"])
+            await refresh_if_needed(session, node, ["current"])
+            await refresh_if_needed(session, node.current, ["required_dimensions"])
 
         instance = cls(session, metric_nodes, use_materialized=use_materialized)
         return instance
@@ -969,7 +977,7 @@ class CubeQueryBuilder:  # pylint: disable=too-many-instance-attributes
         common_parents = group_metrics_by_parent(self.metric_nodes)
         measures_queries = {}
         for parent_node, metrics in common_parents.items():  # type: ignore
-            await self.session.refresh(parent_node, ["current"])
+            await refresh_if_needed(self.session, parent_node, ["current"])
             query_builder = await QueryBuilder.create(self.session, parent_node.current)
             if self._ignore_errors:
                 query_builder = query_builder.ignore_errors()
@@ -999,7 +1007,7 @@ class CubeQueryBuilder:  # pylint: disable=too-many-instance-attributes
                 metric_proj = await self.build_metric_agg(metric_node, parent_node)
                 parent_ast.select.projection.extend(metric_proj)
 
-            await self.session.refresh(parent_node.current, ["columns"])
+            await refresh_if_needed(self.session, parent_node.current, ["columns"])
 
             # Generate semantic types for each
             for expr in parent_ast.select.projection:
@@ -1071,7 +1079,7 @@ class CubeQueryBuilder:  # pylint: disable=too-many-instance-attributes
         """
         Creates an order by ast from the requested order bys
         """
-        temp_orderbys = parse(  # type: ignore
+        temp_orderbys = cached_parse(  # type: ignore
             f"SELECT 1 ORDER BY {','.join(self._orderby)}",
         ).select.organization.order
         valid_sort_items = [
@@ -1117,27 +1125,20 @@ def get_column_from_canonical_dimension(
         column_name = dimension_attr.column_name
 
     # Dimension requested has reference link on node
-    for column in node.columns:
-        if (
-            column.dimension
-            and column.dimension.name == dimension_attr.node_name
-            and column.dimension_column == dimension_attr.column_name
-        ):
-            column_name = column.name
+    dimension_columns = {
+        (col.dimension.name, col.dimension_column): col.name
+        for col in node.columns
+        if col.dimension
+    }
+    key = (dimension_attr.node_name, dimension_attr.column_name)
+    if key in dimension_columns:
+        return dimension_columns[key]
 
     # Dimension referenced was foreign key of dimension link
-    link = next(
-        (
-            link
-            for link in node.dimension_links
-            if dimension_attr.name in link.foreign_keys_reversed
-        ),
-        None,
-    )
-    if link:
-        column_name = FullColumnName(
-            link.foreign_keys_reversed[dimension_attr.name],
-        ).column_name
+    for link in node.dimension_links:
+        foreign_key_column = link.foreign_keys_reversed.get(dimension_attr.name)
+        if foreign_key_column:
+            return FullColumnName(foreign_key_column).column_name
     return column_name
 
 
@@ -1190,7 +1191,7 @@ async def dimension_join_path(
 
     # If it's not a local dimension, traverse the node's dimensions graph
     # This queue tracks the dimension link being processed and the path to that link
-    await session.refresh(node, ["dimension_links"])
+    await refresh_if_needed(session, node, ["dimension_links"])
 
     # Start with first layer of linked dims
     processing_queue = collections.deque(
@@ -1198,11 +1199,16 @@ async def dimension_join_path(
     )
     while processing_queue:
         current_link, join_path = processing_queue.popleft()
-        await session.refresh(current_link, ["dimension"])
+        await refresh_if_needed(session, current_link, ["dimension"])
         if current_link.dimension.name == dimension_attr.node_name:
             return join_path
-        await session.refresh(current_link.dimension, ["current"])
-        await session.refresh(current_link.dimension.current, ["dimension_links"])
+
+        await refresh_if_needed(session, current_link.dimension, ["current"])
+        await refresh_if_needed(
+            session,
+            current_link.dimension.current,
+            ["dimension_links"],
+        )
         processing_queue.extend(
             [
                 (link, join_path + [link])
@@ -1223,8 +1229,12 @@ async def build_dimension_node_query(
     """
     Builds a dimension node query with the requested filters
     """
-    await session.refresh(link.dimension, ["current"])
-    await session.refresh(link.dimension.current, ["availability", "columns"])
+    await refresh_if_needed(session, link.dimension, ["current"])
+    await refresh_if_needed(
+        session,
+        link.dimension.current,
+        ["availability", "columns"],
+    )
     physical_table = get_table_for_node(
         link.dimension.current,
         build_criteria=build_criteria,
@@ -1298,7 +1308,7 @@ async def compile_node_ast(session, node_revision: NodeRevision) -> ast.Query:
     """
     Parses the node's query into an AST and compiles it.
     """
-    node_ast = parse(node_revision.query)
+    node_ast = cached_parse(node_revision.query)
     ctx = CompileContext(session, DJException())
     await node_ast.compile(ctx)
     return node_ast
@@ -1354,8 +1364,8 @@ async def needs_dimension_join(
     """
     if len(join_path) == 1:
         link = join_path[0]
-        await session.refresh(link.dimension, ["current"])
-        await session.refresh(link.dimension.current, ["columns"])
+        await refresh_if_needed(session, link.dimension, ["current"])
+        await refresh_if_needed(session, link.dimension.current, ["columns"])
         if dimension_attribute in link.foreign_keys_reversed:
             return False
     return True
@@ -1430,7 +1440,7 @@ async def build_ast(  # pylint: disable=too-many-arguments,too-many-locals
     This function will apply any filters that can be pushed down to each referenced node's AST
     (filters are only applied if they don't require dimension node joins).
     """
-    await session.refresh(node, ["dimension_links"])
+    await refresh_if_needed(session, node, ["dimension_links"])
     context = CompileContext(session=session, exception=DJException())
     await query.compile(context)
     query.bake_ctes()  # pylint: disable=W0212
@@ -1441,7 +1451,7 @@ async def build_ast(  # pylint: disable=too-many-arguments,too-many-locals
 
     node_to_tables_mapping = get_dj_node_references_from_select(query.select)
     for referenced_node, reference_expressions in node_to_tables_mapping.items():
-        await session.refresh(referenced_node, ["dimension_links"])
+        await refresh_if_needed(session, referenced_node, ["dimension_links"])
 
         for ref_expr in reference_expressions:
 

{datajunction_server-0.0.1a75 → datajunction_server-0.0.1a76.dev0}/datajunction_server/database/dimensionlink.py

@@ -1,4 +1,5 @@
 """Dimension links table."""
+from functools import cached_property
 from typing import TYPE_CHECKING, Dict, List, Optional, Set
 
 from sqlalchemy import JSON, BigInteger, Enum, ForeignKey, Integer
@@ -115,13 +116,14 @@ class DimensionLink(Base):  # pylint: disable=too-few-public-methods
         from datajunction_server.sql.parsing.backends.antlr4 import ast
 
         # Find equality comparions (i.e., fact.order_id = dim.order_id)
+        join_asts = self.joins()
         equality_comparisons = (
             [
                 expr
-                for expr in self.joins()[0].criteria.on.find_all(ast.BinaryOp)  # type: ignore
+                for expr in join_asts[0].criteria.on.find_all(ast.BinaryOp)  # type: ignore
                 if expr.op == ast.BinaryOpKind.Eq
             ]
-            if self.joins()[0].criteria
+            if join_asts[0].criteria
             else []
         )
         mapping = {}
@@ -159,7 +161,7 @@ class DimensionLink(Base):  # pylint: disable=too-few-public-methods
             for fk in self.foreign_keys.keys()
         }
 
-    @hybrid_property
+    @cached_property
     def foreign_keys_reversed(self):
         """
         Returns a mapping from the primary key column(s) on the dimension node to the

{datajunction_server-0.0.1a75 → datajunction_server-0.0.1a76.dev0}/datajunction_server/models/metric.py

@@ -13,6 +13,7 @@ from datajunction_server.models.node import (
     MetricMetadataOutput,
 )
 from datajunction_server.models.query import ColumnMetadata
+from datajunction_server.sql.decompose import Measure, extractor
 from datajunction_server.sql.parsing.backends.antlr4 import ast, parse
 from datajunction_server.transpilation import get_transpilation_plugin
 from datajunction_server.typing import UTCDatetime
@@ -43,6 +44,9 @@ class Metric(BaseModel):
 
     incompatible_druid_functions: List[str]
 
+    measures: List[Measure]
+    derived_sql: str
+
     @classmethod
     def parse_node(cls, node: Node, dims: List[DimensionAttributeOutput]) -> "Metric":
         """
@@ -55,6 +59,7 @@ class Metric(BaseModel):
             for func in functions
             if Dialect.DRUID not in func.dialects
         ]
+        measures, derived_sql = extractor.extract_measures(node.current.query)
         return cls(
             id=node.id,
             name=node.name,
@@ -70,6 +75,8 @@ class Metric(BaseModel):
             metric_metadata=node.current.metric_metadata,
             required_dimensions=[dim.name for dim in node.current.required_dimensions],
             incompatible_druid_functions=incompatible_druid_functions,
+            measures=measures,
+            derived_sql=str(derived_sql).strip(),
         )
 
 

datajunction_server-0.0.1a76.dev0/datajunction_server/sql/decompose.py

@@ -0,0 +1,172 @@
+"""Used for extracting measures form metric definitions."""
+from pydantic import BaseModel
+
+from datajunction_server.enum import StrEnum
+from datajunction_server.sql import functions as dj_functions
+from datajunction_server.sql.parsing.backends.antlr4 import ast, parse
+
+
+class Aggregability(StrEnum):
+    """
+    Type of allowed aggregation for a given measure.
+    """
+
+    FULL = "full"
+    LIMITED = "limited"
+    NONE = "none"
+
+
+class AggregationRule(BaseModel):
+    """
+    The aggregation rule for the measure. If the Aggregability type is LIMITED, the `level` should
+    be specified to highlight the level at which the measure needs to be aggregated in order to
+    support the specified aggregation function.
+
+    For example, consider a metric like COUNT(DISTINCT user_id). It can be decomposed into a
+    single measure with LIMITED aggregability, i.e., it is only aggregatable if the measure is
+    calculated at the `user_id` level:
+    - name: num_users
+      expression: DISTINCT user_id
+      aggregation: COUNT
+      rule:
+        type: LIMITED
+        level: ["user_id"]
+    """
+
+    type: Aggregability = Aggregability.NONE
+    level: list[str] | None = None
+
+
+class Measure(BaseModel):
+    """
+    Measures are aggregated facts (e.g. SUM(view_secs)). They can be optionally combined
+    to build derived metrics, e.g. SUM(clicks) / SUM(view_secs). Combining is optional because
+    a stand-alone measure can itself be a metric.
+    """
+
+    name: str
+    expression: str  # A SQL expression for defining the measure
+    aggregation: str
+    rule: AggregationRule
+
+
+class MeasureExtractor:
+    """
+    Extracts aggregatable measures from a metric definition and generates SQL
+    derived from those measures.
+    """
+
+    def __init__(self):
+        """Register handlers for aggregation functions"""
+        self.handlers = {
+            dj_functions.Sum: self._simple_associative_agg,
+            dj_functions.Count: self._simple_associative_agg,
+            dj_functions.Max: self._simple_associative_agg,
+            dj_functions.Min: self._simple_associative_agg,
+            dj_functions.Avg: self._avg,
+        }
+
+    def extract_measures(self, metric_query: str) -> tuple[list[Measure], ast.Query]:
+        """
+        Decomposes the metric query into its constituent aggregatable measures and
+        constructs a SQL query derived from those measures.
+        """
+        query_ast = parse(metric_query)
+        measures = []
+
+        for idx, func in enumerate(query_ast.find_all(ast.Function)):
+            dj_function = func.function()
+            handler = self.handlers.get(dj_function)
+            if handler and dj_function.is_aggregation:
+                if func_measures := handler(func, idx):  # pragma: no cover
+                    MeasureExtractor.update_ast(func, func_measures)
+                measures.extend(func_measures)
+
+        return measures, query_ast
+
+    def _simple_associative_agg(self, func, idx) -> list[Measure]:
+        """
+        Handles measures decomposition for a single-argument associative aggregation function.
+        Examples: SUM, MAX, MIN, COUNT
+        """
+        arg = func.args[0]
+        measure_name = "_".join(
+            [str(col) for col in arg.find_all(ast.Column)] + [func.name.name.lower()],
+        )
+        return [
+            Measure(
+                name=f"{measure_name}_{idx}",
+                expression=f"{func.quantifier} {arg}" if func.quantifier else str(arg),
+                aggregation=func.name.name.upper(),
+                rule=AggregationRule(
+                    type=Aggregability.FULL
+                    if func.quantifier != ast.SetQuantifier.Distinct
+                    else Aggregability.LIMITED,
+                ),
+            ),
+        ]
+
+    def _avg(self, func, idx) -> list[Measure]:
+        """
+        Handles measures decomposition for AVG (it requires both the SUM and COUNT
+        of the selected measure).
+        """
+        arg = func.args[0]
+        measure_name = "_".join(
+            [str(col) for col in arg.find_all(ast.Column)]
+            + [dj_functions.Sum.__name__.lower()],
+        )
+        return [
+            Measure(
+                name=f"{measure_name}_{idx}",
+                expression=str(arg),
+                aggregation=dj_functions.Sum.__name__.upper(),
+                rule=AggregationRule(
+                    type=Aggregability.FULL
+                    if func.quantifier != ast.SetQuantifier.Distinct
+                    else Aggregability.LIMITED,
+                ),
+            ),
+            Measure(
+                name="count",
+                expression="1",
+                aggregation=dj_functions.Count.__name__.upper(),
+                rule=AggregationRule(
+                    type=Aggregability.FULL
+                    if func.quantifier != ast.SetQuantifier.Distinct
+                    else Aggregability.LIMITED,
+                ),
+            ),
+        ]
+
+    @staticmethod
+    def update_ast(func, measures: list[Measure]):
+        """
+        Updates the query AST based on the measures derived from the function.
+        """
+        if func.function() == dj_functions.Avg:
+            func.parent.replace(
+                from_=func,
+                to=ast.BinaryOp(
+                    op=ast.BinaryOpKind.Divide,
+                    left=ast.Function(
+                        ast.Name("SUM"),
+                        args=[ast.Column(ast.Name(measures[0].name))],
+                    ),
+                    right=ast.Function(
+                        ast.Name("COUNT"),
+                        args=[ast.Column(ast.Name(measures[1].name))],
+                    ),
+                ),
+            )
+        elif (
+            func.function() == dj_functions.Count
+            and func.quantifier != ast.SetQuantifier.Distinct
+        ):
+            func.name.name = "SUM"
+            func.args = [ast.Column(ast.Name(measure.name)) for measure in measures]
+        else:
+            func.args = [ast.Column(ast.Name(measure.name)) for measure in measures]
+
+
+extractor = MeasureExtractor()

{datajunction_server-0.0.1a75 → datajunction_server-0.0.1a76.dev0}/datajunction_server/sql/parsing/ast.py

@@ -1679,6 +1679,19 @@ class Over(Expression):
         return f"OVER ({consolidated_by}{window_frame})"
 
 
+# pylint: disable=C0103
+class SetQuantifier(DJEnum):
+    """
+    The accepted set quantifiers
+    """
+
+    All = "ALL"
+    Distinct = "DISTINCT"
+
+    def __str__(self):
+        return self.value
+
+
 @dataclass(eq=False)
 class Function(Named, Operation):
     """
@@ -1686,7 +1699,7 @@ class Function(Named, Operation):
     """
 
     args: List[Expression] = field(default_factory=list)
-    quantifier: str = ""
+    quantifier: SetQuantifier | None = None
     over: Optional[Over] = None
     args_compiled: bool = False
 

{datajunction_server-0.0.1a75 → datajunction_server-0.0.1a76.dev0}/datajunction_server/sql/parsing/backends/antlr4.py

@@ -1,7 +1,9 @@
 # pylint: skip-file
 # mypy: ignore-errors
+import copy
 import inspect
 import logging
+from functools import lru_cache
 from typing import TYPE_CHECKING, List, Optional, Tuple, Union, cast
 
 import antlr4
@@ -179,6 +181,14 @@ def parse_rule(sql: str, rule: str) -> Union[ast.Node, "ColumnType"]:
     return ast_tree
 
 
+@lru_cache(maxsize=128)
+def _cached_parse(sql: Optional[str]) -> ast.Query:
+    """
+    Parse a string sql query into a DJ query AST and cache it.
+    """
+    return parse(sql)
+
+
 def parse(sql: Optional[str]) -> ast.Query:
     """
     Parse a string sql query into a DJ ast Query
@@ -188,6 +198,13 @@ def parse(sql: Optional[str]) -> ast.Query:
     return cast(ast.Query, parse_rule(sql, "singleStatement"))
 
 
+def cached_parse(sql: Optional[str]) -> ast.Query:
+    """
+    Parse a string sql query into a DJ ast Query
+    """
+    return copy.deepcopy(_cached_parse(sql))
+
+
 TERMINAL_NODE = antlr4.tree.Tree.TerminalNodeImpl
 
 
@@ -659,7 +676,9 @@ def _(ctx: sbp.DereferenceContext):
 @visit.register
 def _(ctx: sbp.FunctionCallContext):
     name = visit(ctx.functionName())
-    quantifier = visit(ctx.setQuantifier()) if ctx.setQuantifier() else ""
+    quantifier = (
+        ast.SetQuantifier(visit(ctx.setQuantifier())) if ctx.setQuantifier() else None
+    )
     over = visit(ctx.windowSpec()) if ctx.windowSpec() else None
     args = visit(ctx.argument)
     return ast.Function(name, args, quantifier=quantifier, over=over)

{datajunction_server-0.0.1a75 → datajunction_server-0.0.1a76.dev0}/datajunction_server/utils.py

@@ -16,6 +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.ext.asyncio import (
     AsyncEngine,
     AsyncSession,
@@ -161,6 +162,22 @@ async def get_session() -> AsyncIterator[AsyncSession]:
         await session.close()
 
 
+async def refresh_if_needed(session: AsyncSession, obj, attributes: list[str]):
+    """
+    Conditionally refresh a list of attributes for a SQLAlchemy ORM object.
+    """
+    attributes_to_refresh = []
+
+    for attr_name in attributes:
+        try:
+            getattr(obj, attr_name)
+        except MissingGreenlet:
+            attributes_to_refresh.append(attr_name)
+
+    if attributes_to_refresh:
+        await session.refresh(obj, attributes_to_refresh)
+
+
 def get_query_service_client() -> Optional[QueryServiceClient]:
     """
     Return query service client