Skip to content

Commit

Permalink
Update other execution plan tasks.
Browse files Browse the repository at this point in the history
  • Loading branch information
plypaul committed Dec 14, 2024
1 parent fd38bcb commit 3c13f7b
Show file tree
Hide file tree
Showing 3 changed files with 14 additions and 14 deletions.
2 changes: 1 addition & 1 deletion metricflow/execution/dataflow_to_execution.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,7 @@ def visit_write_to_result_table_node(self, node: WriteToResultTableNode) -> Conv
leaf_tasks=(
SelectSqlQueryToTableTask.create(
sql_client=self._sql_client,
sql_query=SqlStatement(
sql_statement=SqlStatement(
sql=render_sql_result.sql,
bind_parameter_set=render_sql_result.bind_parameter_set,
),
Expand Down
24 changes: 12 additions & 12 deletions metricflow/execution/execution_plan.py
Original file line number Diff line number Diff line change
Expand Up @@ -167,13 +167,13 @@ class SelectSqlQueryToTableTask(ExecutionPlanTask):
@staticmethod
def create( # noqa: D102
sql_client: SqlClient,
sql_query: SqlStatement,
sql_statement: SqlStatement,
output_table: SqlTable,
parent_nodes: Sequence[ExecutionPlanTask] = (),
) -> SelectSqlQueryToTableTask:
return SelectSqlQueryToTableTask(
sql_client=sql_client,
sql_statement=sql_query,
sql_statement=sql_statement,
output_table=output_table,
parent_nodes=tuple(parent_nodes),
)
Expand All @@ -188,31 +188,31 @@ def description(self) -> str: # noqa: D102

@property
def displayed_properties(self) -> Sequence[DisplayedProperty]: # noqa: D102
sql_query = self.sql_statement
assert sql_query is not None, f"{self.sql_statement=} should have been set during creation."
sql_statement = self.sql_statement
assert sql_statement 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="sql_statement", value=sql_statement.sql),
DisplayedProperty(key="output_table", value=self.output_table),
DisplayedProperty(key="bind_parameter_set", value=sql_query.bind_parameter_set),
DisplayedProperty(key="bind_parameter_set", value=sql_statement.bind_parameter_set),
)

def execute(self) -> TaskExecutionResult: # noqa: D102
sql_query = self.sql_statement
assert sql_query is not None, f"{self.sql_statement=} should have been set during creation."
sql_statement = self.sql_statement
assert sql_statement 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}")
logger.debug(LazyFormat(lambda: f"Creating table {self.output_table} using a query"))
self.sql_client.execute(
sql_query.sql,
sql_bind_parameter_set=sql_query.bind_parameter_set,
sql_statement.sql,
sql_bind_parameter_set=sql_statement.bind_parameter_set,
)

end_time = time.time()
return TaskExecutionResult(start_time=start_time, end_time=end_time, sql=sql_query.sql)
return TaskExecutionResult(start_time=start_time, end_time=end_time, sql=sql_statement.sql)

def __repr__(self) -> str: # noqa: D105
return f"{self.__class__.__name__}(sql_query='{self.sql_statement}', output_table={self.output_table})"
return f"{self.__class__.__name__}(sql_statement={self.sql_statement!r}', output_table={self.output_table})"


class ExecutionPlan(MetricFlowDag[ExecutionPlanTask]):
Expand Down
2 changes: 1 addition & 1 deletion tests_metricflow/execution/test_tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ def test_write_table_task( # noqa: D103
output_table = SqlTable(schema_name=mf_test_configuration.mf_system_schema, table_name=f"test_table_{random_id()}")
task = SelectSqlQueryToTableTask.create(
sql_client=sql_client,
sql_query=SqlStatement(
sql_statement=SqlStatement(
sql=f"CREATE TABLE {output_table.sql} AS SELECT 1 AS foo",
bind_parameter_set=SqlBindParameterSet(),
),
Expand Down

0 comments on commit 3c13f7b

Please sign in to comment.