From 55896cd6a28193fb00fca7fefc18c0d8710b5b74 Mon Sep 17 00:00:00 2001 From: Courtney Holcomb Date: Mon, 25 Sep 2023 13:31:49 -0700 Subject: [PATCH] Alias unions & update docstring --- metricflow/engine/metricflow_engine.py | 23 ++++++++----------- metricflow/protocols/query_parameter.py | 8 +++++-- metricflow/query/query_parser.py | 10 ++++---- .../specs/query_param_implementations.py | 12 ++++------ .../test/integration/test_configured_cases.py | 8 +++---- metricflow/test/query/test_query_parser.py | 6 ++--- 6 files changed, 32 insertions(+), 35 deletions(-) diff --git a/metricflow/engine/metricflow_engine.py b/metricflow/engine/metricflow_engine.py index a42718f3a8..bf66a994b6 100644 --- a/metricflow/engine/metricflow_engine.py +++ b/metricflow/engine/metricflow_engine.py @@ -5,7 +5,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from enum import Enum -from typing import List, Optional, Sequence, Tuple, Union +from typing import List, Optional, Sequence, Tuple import pandas as pd from dbt_semantic_interfaces.implementations.elements.dimension import PydanticDimensionTypeParams @@ -47,12 +47,7 @@ DataflowToExecutionPlanConverter, ) from metricflow.plan_conversion.dataflow_to_sql import DataflowToSqlQueryPlanConverter -from metricflow.protocols.query_parameter import ( - GroupByQueryParameter, - MetricQueryParameter, - OrderByQueryParameter, - TimeDimensionQueryParameter, -) +from metricflow.protocols.query_parameter import GroupByParameter, MetricQueryParameter, OrderByQueryParameter from metricflow.protocols.sql_client import SqlClient from metricflow.query.query_exceptions import InvalidQueryException from metricflow.query.query_parser import MetricFlowQueryParser @@ -90,13 +85,15 @@ class MetricFlowQueryRequest: """Encapsulates the parameters for a metric query. metric_names: Names of the metrics to query. + metrics: Metric objects to query. group_by_names: Names of the dimensions and entities to query. + group_by: Dimension or entity objects to query. limit: Limit the result to this many rows. time_constraint_start: Get data for the start of this time range. time_constraint_end: Get data for the end of this time range. where_constraint: A SQL string using group by names that can be used like a where clause on the output data. - order_by_names: metric and group by names to order by. A "-" can be used to specify reverse order e.g. "-ds" - order_by: metric and group by objects to order by + order_by_names: metric and group by names to order by. A "-" can be used to specify reverse order e.g. "-ds". + order_by: metric, dimension, or entity objects to order by. output_table: If specified, output the result data to this table instead of a result dataframe. sql_optimization_level: The level of optimization for the generated SQL. query_type: Type of MetricFlow query. @@ -106,7 +103,7 @@ class MetricFlowQueryRequest: metric_names: Optional[Sequence[str]] = None metrics: Optional[Sequence[MetricQueryParameter]] = None group_by_names: Optional[Sequence[str]] = None - group_by: Optional[Tuple[Union[GroupByQueryParameter, TimeDimensionQueryParameter], ...]] = None + group_by: Optional[Tuple[GroupByParameter, ...]] = None limit: Optional[int] = None time_constraint_start: Optional[datetime.datetime] = None time_constraint_end: Optional[datetime.datetime] = None @@ -122,7 +119,7 @@ def create_with_random_request_id( # noqa: D metric_names: Optional[Sequence[str]] = None, metrics: Optional[Sequence[MetricQueryParameter]] = None, group_by_names: Optional[Sequence[str]] = None, - group_by: Optional[Tuple[Union[GroupByQueryParameter, TimeDimensionQueryParameter], ...]] = None, + group_by: Optional[Tuple[GroupByParameter, ...]] = None, limit: Optional[int] = None, time_constraint_start: Optional[datetime.datetime] = None, time_constraint_end: Optional[datetime.datetime] = None, @@ -294,7 +291,7 @@ def explain_get_dimension_values( # noqa: D metric_names: Optional[List[str]] = None, metrics: Optional[Sequence[MetricQueryParameter]] = None, get_group_by_values: Optional[str] = None, - group_by: Optional[Union[GroupByQueryParameter, TimeDimensionQueryParameter]] = None, + group_by: Optional[GroupByParameter] = None, time_constraint_start: Optional[datetime.datetime] = None, time_constraint_end: Optional[datetime.datetime] = None, ) -> MetricFlowExplainResult: @@ -691,7 +688,7 @@ def explain_get_dimension_values( # noqa: D metric_names: Optional[List[str]] = None, metrics: Optional[Sequence[MetricQueryParameter]] = None, get_group_by_values: Optional[str] = None, - group_by: Optional[Union[GroupByQueryParameter, TimeDimensionQueryParameter]] = None, + group_by: Optional[GroupByParameter] = None, time_constraint_start: Optional[datetime.datetime] = None, time_constraint_end: Optional[datetime.datetime] = None, ) -> MetricFlowExplainResult: diff --git a/metricflow/protocols/query_parameter.py b/metricflow/protocols/query_parameter.py index f2e27ec393..79e5db93c7 100644 --- a/metricflow/protocols/query_parameter.py +++ b/metricflow/protocols/query_parameter.py @@ -18,7 +18,7 @@ def name(self) -> str: @runtime_checkable -class GroupByQueryParameter(Protocol): +class DimensionOrEntityQueryParameter(Protocol): """Generic group by parameter for queries. Might be an entity or a dimension.""" @property @@ -45,11 +45,15 @@ def date_part(self) -> Optional[DatePart]: raise NotImplementedError +GroupByParameter = Union[DimensionOrEntityQueryParameter, TimeDimensionQueryParameter] +InputOrderByParameter = Union[MetricQueryParameter, GroupByParameter] + + class OrderByQueryParameter(Protocol): """Parameter to order by, specifying ascending or descending.""" @property - def order_by(self) -> Union[MetricQueryParameter, GroupByQueryParameter, TimeDimensionQueryParameter]: + def order_by(self) -> InputOrderByParameter: """Parameter to order results by.""" raise NotImplementedError diff --git a/metricflow/query/query_parser.py b/metricflow/query/query_parser.py index 701c059a3c..4312b44764 100644 --- a/metricflow/query/query_parser.py +++ b/metricflow/query/query_parser.py @@ -4,7 +4,7 @@ import logging import time from dataclasses import dataclass -from typing import Dict, List, Optional, Sequence, Tuple, Union +from typing import Dict, List, Optional, Sequence, Tuple from dbt_semantic_interfaces.call_parameter_sets import ParseWhereFilterException from dbt_semantic_interfaces.implementations.filters.where_filter import PydanticWhereFilter @@ -29,7 +29,7 @@ from metricflow.model.semantic_manifest_lookup import SemanticManifestLookup from metricflow.naming.linkable_spec_name import StructuredLinkableSpecName from metricflow.protocols.query_parameter import ( - GroupByQueryParameter, + GroupByParameter, MetricQueryParameter, OrderByQueryParameter, TimeDimensionQueryParameter, @@ -177,7 +177,7 @@ def parse_and_validate_query( metric_names: Optional[Sequence[str]] = None, metrics: Optional[Sequence[MetricQueryParameter]] = None, group_by_names: Optional[Sequence[str]] = None, - group_by: Optional[Tuple[Union[GroupByQueryParameter, TimeDimensionQueryParameter], ...]] = None, + group_by: Optional[Tuple[GroupByParameter, ...]] = None, limit: Optional[int] = None, time_constraint_start: Optional[datetime.datetime] = None, time_constraint_end: Optional[datetime.datetime] = None, @@ -316,7 +316,7 @@ def _parse_and_validate_query( metric_names: Optional[Sequence[str]] = None, metrics: Optional[Sequence[MetricQueryParameter]] = None, group_by_names: Optional[Sequence[str]] = None, - group_by: Optional[Tuple[Union[GroupByQueryParameter, TimeDimensionQueryParameter], ...]] = None, + group_by: Optional[Tuple[GroupByParameter, ...]] = None, limit: Optional[int] = None, time_constraint_start: Optional[datetime.datetime] = None, time_constraint_end: Optional[datetime.datetime] = None, @@ -665,7 +665,7 @@ def _parse_group_by( self, metric_references: Sequence[MetricReference], group_by_names: Optional[Sequence[str]] = None, - group_by: Optional[Tuple[Union[GroupByQueryParameter, TimeDimensionQueryParameter], ...]] = None, + group_by: Optional[Tuple[GroupByParameter, ...]] = None, ) -> QueryTimeLinkableSpecSet: """Convert the linkable spec names into the respective specification objects.""" # TODO: refactor to only support group_by object inputs (removing group_by_names param) diff --git a/metricflow/specs/query_param_implementations.py b/metricflow/specs/query_param_implementations.py index dd534eace2..5b06b13f14 100644 --- a/metricflow/specs/query_param_implementations.py +++ b/metricflow/specs/query_param_implementations.py @@ -1,16 +1,12 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Optional, Union +from typing import Optional from dbt_semantic_interfaces.type_enums.time_granularity import TimeGranularity from metricflow.naming.linkable_spec_name import StructuredLinkableSpecName -from metricflow.protocols.query_parameter import ( - GroupByQueryParameter, - MetricQueryParameter, - TimeDimensionQueryParameter, -) +from metricflow.protocols.query_parameter import InputOrderByParameter from metricflow.time.date_part import DatePart @@ -29,7 +25,7 @@ def __post_init__(self) -> None: # noqa: D @dataclass(frozen=True) -class GroupByParameter: +class DimensionOrEntityParameter: """Group by parameter requested in a query. Might represent an entity or a dimension. @@ -49,5 +45,5 @@ class MetricParameter: class OrderByParameter: """Order by requested in a query.""" - order_by: Union[MetricQueryParameter, GroupByQueryParameter, TimeDimensionQueryParameter] + order_by: InputOrderByParameter descending: bool = False diff --git a/metricflow/test/integration/test_configured_cases.py b/metricflow/test/integration/test_configured_cases.py index 366d8e5aa7..eec106d555 100644 --- a/metricflow/test/integration/test_configured_cases.py +++ b/metricflow/test/integration/test_configured_cases.py @@ -18,9 +18,9 @@ from metricflow.plan_conversion.column_resolver import ( DunderColumnAssociationResolver, ) -from metricflow.protocols.query_parameter import GroupByQueryParameter +from metricflow.protocols.query_parameter import DimensionOrEntityQueryParameter from metricflow.protocols.sql_client import SqlClient -from metricflow.specs.query_param_implementations import GroupByParameter, TimeDimensionParameter +from metricflow.specs.query_param_implementations import DimensionOrEntityParameter, TimeDimensionParameter from metricflow.sql.sql_exprs import ( SqlCastToTimestampExpression, SqlColumnReference, @@ -256,7 +256,7 @@ def test_case( check_query_helpers = CheckQueryHelpers(sql_client) - group_by: List[GroupByQueryParameter] = [] + group_by: List[DimensionOrEntityQueryParameter] = [] for group_by_kwargs in case.group_by_objs: kwargs = copy(group_by_kwargs) date_part = kwargs.get("date_part") @@ -268,7 +268,7 @@ def test_case( kwargs["grain"] = TimeGranularity(grain) group_by.append(TimeDimensionParameter(**kwargs)) else: - group_by.append(GroupByParameter(**kwargs)) + group_by.append(DimensionOrEntityParameter(**kwargs)) query_result = engine.query( MetricFlowQueryRequest.create_with_random_request_id( metric_names=case.metrics, diff --git a/metricflow/test/query/test_query_parser.py b/metricflow/test/query/test_query_parser.py index a5644406c2..51f272879c 100644 --- a/metricflow/test/query/test_query_parser.py +++ b/metricflow/test/query/test_query_parser.py @@ -15,7 +15,7 @@ from metricflow.query.query_exceptions import InvalidQueryException from metricflow.query.query_parser import MetricFlowQueryParser from metricflow.specs.query_param_implementations import ( - GroupByParameter, + DimensionOrEntityParameter, MetricParameter, OrderByParameter, TimeDimensionParameter, @@ -196,8 +196,8 @@ def test_query_parser_with_object_params(bookings_query_parser: MetricFlowQueryP Metric = namedtuple("Metric", ["name", "descending"]) metric = Metric("bookings", False) group_by = ( - GroupByParameter("booking__is_instant"), - GroupByParameter("listing"), + DimensionOrEntityParameter("booking__is_instant"), + DimensionOrEntityParameter("listing"), TimeDimensionParameter(MTD), ) order_by = (