From b4082afbd36f67c371853aadeaf986535f45c794 Mon Sep 17 00:00:00 2001 From: Paul Yang Date: Tue, 10 Dec 2024 18:47:21 -0800 Subject: [PATCH] Rename `ExecutionPlanTask.sql_query`. --- metricflow/engine/metricflow_engine.py | 2 +- metricflow/execution/execution_plan.py | 32 ++++++++++++------------- tests_metricflow/execution/noop_task.py | 2 +- 3 files changed, 18 insertions(+), 18 deletions(-) diff --git a/metricflow/engine/metricflow_engine.py b/metricflow/engine/metricflow_engine.py index 4d6319eca5..f494a4b444 100644 --- a/metricflow/engine/metricflow_engine.py +++ b/metricflow/engine/metricflow_engine.py @@ -185,7 +185,7 @@ def sql_statement(self) -> SqlStatement: f"Multiple tasks in the execution plan not yet supported. Got tasks: {execution_plan.tasks}" ) - sql_query = execution_plan.tasks[0].sql_query + sql_query = execution_plan.tasks[0].sql_statement if not sql_query: raise NotImplementedError( f"Execution plan tasks without a SQL query not yet supported. Got tasks: {execution_plan.tasks}" diff --git a/metricflow/execution/execution_plan.py b/metricflow/execution/execution_plan.py index a29f956f2b..dd684b3f82 100644 --- a/metricflow/execution/execution_plan.py +++ b/metricflow/execution/execution_plan.py @@ -27,10 +27,10 @@ class ExecutionPlanTask(DagNode["ExecutionPlanTask"], Visitable, ABC): for these nodes as it seems more intuitive. Attributes: - sql_query: If this runs a SQL query, return the associated SQL. + sql_statement: If this runs a SQL query, return the associated SQL. """ - sql_query: Optional[SqlStatement] + sql_statement: Optional[SqlStatement] @abstractmethod def execute(self) -> TaskExecutionResult: @@ -93,7 +93,7 @@ class SelectSqlQueryToDataTableTask(ExecutionPlanTask): Attributes: sql_client: The SQL client used to run the query. - sql_query: The SQL query to run. + sql_statement: The SQL query to run. parent_nodes: The parent tasks for this execution plan task. """ @@ -108,7 +108,7 @@ def create( # noqa: D102 ) -> SelectSqlQueryToDataTableTask: return SelectSqlQueryToDataTableTask( sql_client=sql_client, - sql_query=sql_query, + sql_statement=sql_query, parent_nodes=tuple(parent_nodes), ) @@ -122,14 +122,14 @@ def description(self) -> str: # noqa: D102 @property def displayed_properties(self) -> Sequence[DisplayedProperty]: # noqa: D102 - sql_query = self.sql_query - assert sql_query is not None, f"{self.sql_query=} should have been set during creation." + sql_query = self.sql_statement + assert sql_query is not None, f"{self.sql_statement=} should have been set during creation." return tuple(super().displayed_properties) + (DisplayedProperty(key="sql_query", value=sql_query.sql),) def execute(self) -> TaskExecutionResult: # noqa: D102 start_time = time.time() - sql_query = self.sql_query - assert sql_query is not None, f"{self.sql_query=} should have been set during creation." + sql_query = self.sql_statement + assert sql_query is not None, f"{self.sql_statement=} should have been set during creation." df = self.sql_client.query( sql_query.sql, @@ -146,7 +146,7 @@ def execute(self) -> TaskExecutionResult: # noqa: D102 ) def __repr__(self) -> str: # noqa: D105 - return f"{self.__class__.__name__}(sql_query='{self.sql_query}')" + return f"{self.__class__.__name__}(sql_query='{self.sql_statement}')" @dataclass(frozen=True) @@ -157,7 +157,7 @@ class SelectSqlQueryToTableTask(ExecutionPlanTask): Attributes: sql_client: The SQL client used to run the query. - sql_query: The SQL query to run. + sql_statement: The SQL query to run. output_table: The table where the results will be written. """ @@ -173,7 +173,7 @@ def create( # noqa: D102 ) -> SelectSqlQueryToTableTask: return SelectSqlQueryToTableTask( sql_client=sql_client, - sql_query=sql_query, + sql_statement=sql_query, output_table=output_table, parent_nodes=tuple(parent_nodes), ) @@ -188,8 +188,8 @@ def description(self) -> str: # noqa: D102 @property def displayed_properties(self) -> Sequence[DisplayedProperty]: # noqa: D102 - sql_query = self.sql_query - assert sql_query is not None, f"{self.sql_query=} should have been set during creation." + sql_query = self.sql_statement + assert sql_query is not None, f"{self.sql_statement=} should have been set during creation." return tuple(super().displayed_properties) + ( DisplayedProperty(key="sql_query", value=sql_query.sql), DisplayedProperty(key="output_table", value=self.output_table), @@ -197,8 +197,8 @@ def displayed_properties(self) -> Sequence[DisplayedProperty]: # noqa: D102 ) def execute(self) -> TaskExecutionResult: # noqa: D102 - sql_query = self.sql_query - assert sql_query is not None, f"{self.sql_query=} should have been set during creation." + sql_query = self.sql_statement + assert sql_query is not None, f"{self.sql_statement=} should have been set during creation." start_time = time.time() logger.debug(LazyFormat(lambda: f"Dropping table {self.output_table} in case it already exists")) self.sql_client.execute(f"DROP TABLE IF EXISTS {self.output_table.sql}") @@ -212,7 +212,7 @@ def execute(self) -> TaskExecutionResult: # noqa: D102 return TaskExecutionResult(start_time=start_time, end_time=end_time, sql=sql_query.sql) def __repr__(self) -> str: # noqa: D105 - return f"{self.__class__.__name__}(sql_query='{self.sql_query}', output_table={self.output_table})" + return f"{self.__class__.__name__}(sql_query='{self.sql_statement}', output_table={self.output_table})" class ExecutionPlan(MetricFlowDag[ExecutionPlanTask]): diff --git a/tests_metricflow/execution/noop_task.py b/tests_metricflow/execution/noop_task.py index 8f65ec049a..ad06af4966 100644 --- a/tests_metricflow/execution/noop_task.py +++ b/tests_metricflow/execution/noop_task.py @@ -35,7 +35,7 @@ def create( # noqa: D102 ) -> NoOpExecutionPlanTask: return NoOpExecutionPlanTask( parent_nodes=tuple(parent_tasks), - sql_query=None, + sql_statement=None, should_error=should_error, )