You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "pankajkoti (via GitHub)" <gi...@apache.org> on 2025/03/28 19:20:59 UTC
[PR] Add DatabricksSQLStatementsOperator [airflow]
pankajkoti opened a new pull request, #48507:
URL: https://github.com/apache/airflow/pull/48507
<!--
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements. See the NOTICE file
distributed with this work for additional information
regarding copyright ownership. The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied. See the License for the
specific language governing permissions and limitations
under the License.
-->
<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.
Feel free to ping committers for the review!
In case of an existing issue, reference it using one of the following:
closes: #ISSUE
related: #ISSUE
How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
**WIP**
Pending: tests
Verified that the operator works as expected by running an Example DAG with Airflow 2.10.5 and confirming that the deferred execution functions properly.
<!-- Please keep an empty line above the dashes. -->
---
**^ Add meaningful description above**
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information.
In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in a newsfragment file, named `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add DatabricksSQLStatementsOperator [airflow]
Posted by "pankajkoti (via GitHub)" <gi...@apache.org>.
pankajkoti commented on code in PR #48507:
URL: https://github.com/apache/airflow/pull/48507#discussion_r2025311741
##########
providers/databricks/src/airflow/providers/databricks/operators/databricks.py:
##########
@@ -969,6 +978,183 @@ def on_kill(self) -> None:
self.log.error("Error: Task: %s with invalid run_id was requested to be cancelled.", self.task_id)
+class DatabricksSQLStatementsOperator(BaseOperator):
+ """
+ Submits a Databricks SQL Statement to Databricks using the api/2.0/sql/statements/ API endpoint.
+
+ See: https://docs.databricks.com/api/workspace/statementexecution
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:DatabricksSQLStatementsOperator`
+
+ :param statement: The SQL statement to execute. The statement can optionally be parameterized, see parameters.
+ :param warehouse_id: Warehouse upon which to execute a statement.
+ :param catalog: Sets default catalog for statement execution, similar to USE CATALOG in SQL.
+ :param schema: Sets default schema for statement execution, similar to USE SCHEMA in SQL.
+ :param parameters: A list of parameters to pass into a SQL statement containing parameter markers.
+
+ .. seealso::
+ https://docs.databricks.com/api/workspace/statementexecution/executestatement#parameters
+ :param wait_for_termination: if we should wait for termination of the statement execution. ``True`` by default.
+ :param databricks_conn_id: Reference to the :ref:`Databricks connection <howto/connection:databricks>`.
+ By default and in the common case this will be ``databricks_default``. To use
+ token based authentication, provide the key ``token`` in the extra field for the
+ connection and create the key ``host`` and leave the ``host`` field empty. (templated)
+ :param polling_period_seconds: Controls the rate which we poll for the result of
+ this statement. By default the operator will poll every 30 seconds.
+ :param databricks_retry_limit: Amount of times retry if the Databricks backend is
+ unreachable. Its value must be greater than or equal to 1.
+ :param databricks_retry_delay: Number of seconds to wait between retries (it
+ might be a floating point number).
+ :param databricks_retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` class.
+ :param do_xcom_push: Whether we should push statement_id to xcom.
+ :param deferrable: Run operator in the deferrable mode.
+ """
+
+ # Used in airflow.models.BaseOperator
+ template_fields: Sequence[str] = ("databricks_conn_id",)
+ template_ext: Sequence[str] = (".json-tpl",)
+ # Databricks brand color (blue) under white text
+ ui_color = "#1CB1C2"
+ ui_fgcolor = "#fff"
+
+ def __init__(
+ self,
+ statement: str,
+ warehouse_id: str,
+ *,
+ catalog: str | None = None,
+ schema: str | None = None,
+ parameters: list[dict[str, Any]] | None = None,
+ databricks_conn_id: str = "databricks_default",
+ polling_period_seconds: int = 30,
+ databricks_retry_limit: int = 3,
+ databricks_retry_delay: int = 1,
+ databricks_retry_args: dict[Any, Any] | None = None,
+ do_xcom_push: bool = True,
+ wait_for_termination: bool = True,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ) -> None:
+ """Create a new ``DatabricksSubmitRunOperator``."""
+ super().__init__(**kwargs)
+ self.statement = statement
+ self.warehouse_id = warehouse_id
+ self.catalog = catalog
+ self.schema = schema
+ self.parameters = parameters
+ self.databricks_conn_id = databricks_conn_id
+ self.polling_period_seconds = polling_period_seconds
+ self.databricks_retry_limit = databricks_retry_limit
+ self.databricks_retry_delay = databricks_retry_delay
+ self.databricks_retry_args = databricks_retry_args
+ self.wait_for_termination = wait_for_termination
+ self.deferrable = deferrable
+
+ # This variable will be used in case our task gets killed.
+ self.statement_id: str | None = None
+ self.do_xcom_push = do_xcom_push
+
+ @cached_property
+ def _hook(self):
+ return self._get_hook(caller="DatabricksSQLStatementsOperator")
+
+ def _get_hook(self, caller: str) -> DatabricksHook:
+ return DatabricksHook(
+ self.databricks_conn_id,
+ retry_limit=self.databricks_retry_limit,
+ retry_delay=self.databricks_retry_delay,
+ retry_args=self.databricks_retry_args,
+ caller=caller,
+ )
+
+ def _handle_operator_execution(self) -> None:
+ while True:
+ statement_state = self._hook.get_sql_statement_state(self.statement_id)
+ if statement_state.is_terminal:
+ if statement_state.is_successful:
+ self.log.info("%s completed successfully.", self.task_id)
+ return
+ error_message = (
+ f"{self.task_id} failed with terminal state: {statement_state.state} "
+ f"and with the error code {statement_state.error_code} "
+ f"and error message {statement_state.error_message}"
+ )
+ raise AirflowException(error_message)
+
+ self.log.info("%s in run state: %s", self.task_id, statement_state.state)
+ self.log.info("Sleeping for %s seconds.", self.polling_period_seconds)
+ time.sleep(self.polling_period_seconds)
+
+ def _handle_deferrable_operator_execution(self) -> None:
+ statement_state = self._hook.get_sql_statement_state(self.statement_id)
+ if not statement_state.is_terminal:
+ self.defer(
+ trigger=DatabricksSQLStatementExecutionTrigger(
+ statement_id=self.statement_id,
+ databricks_conn_id=self.databricks_conn_id,
+ polling_period_seconds=self.polling_period_seconds,
+ retry_limit=self.databricks_retry_limit,
+ retry_delay=self.databricks_retry_delay,
+ retry_args=self.databricks_retry_args,
+ ),
+ method_name=DEFER_METHOD_NAME,
Review Comment:
I observed that the module had a constant defined for this https://github.com/apache/airflow/blob/e85b61dd40ef6a1a16fd9260aff4fe6f228be6e8/providers/databricks/src/airflow/providers/databricks/operators/databricks.py#L58 whose value is basically execute_complete. I am simply re-using that constant here :)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add DatabricksSQLStatementsOperator [airflow]
Posted by "rawwar (via GitHub)" <gi...@apache.org>.
rawwar commented on code in PR #48507:
URL: https://github.com/apache/airflow/pull/48507#discussion_r2024139221
##########
providers/databricks/src/airflow/providers/databricks/hooks/databricks.py:
##########
@@ -189,6 +189,65 @@ def from_json(cls, data: str) -> ClusterState:
return ClusterState(**json.loads(data))
+class SQLStatementState:
+ """Utility class for the SQL statement state concept of Databricks statements."""
+
+ SQL_STATEMENT_LIFE_CYCLE_STATES = [
+ "PENDING",
+ "RUNNING",
+ "SUCCEEDED",
+ "FAILED",
+ "CANCELED",
+ "CLOSED",
+ ]
+
+ def __init__(
+ self, state: str = "", error_code: str = "", error_message: str = "", *args, **kwargs
+ ) -> None:
+ if state not in self.SQL_STATEMENT_LIFE_CYCLE_STATES:
+ raise AirflowException(
+ f"Unexpected SQL statement life cycle state: {state}: If the state has "
+ "been introduced recently, please check the Databricks user "
+ "guide for troubleshooting information"
+ )
+
+ self.state = state
+ self.error_code = error_code
+ self.error_message = error_message
+
+ @property
+ def is_terminal(self) -> bool:
+ """True if the current state is a terminal state."""
+ return self.state in ("SUCCEEDED", "FAILED", "CANCELED", "CLOSED")
+
+ @property
+ def is_running(self) -> bool:
+ """True if the current state is running."""
+ return self.state in ("PENDING", "RUNNING")
+
+ @property
+ def is_successful(self) -> bool:
+ """True if the state is SUCCEEDED."""
+ return self.state == "SUCCEEDED"
+
+ def __eq__(self, other) -> bool:
+ return (
+ self.state == other.state
+ and self.error_code == other.error_code
+ and self.error_message == other.error_message
+ )
Review Comment:
I noticed this in RunState and made sense, given, we don't want to compare between other type of state objects. Maybe, easier to just make it `other: SQLStatementState`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add DatabricksSQLStatementsOperator [airflow]
Posted by "pankajkoti (via GitHub)" <gi...@apache.org>.
pankajkoti commented on code in PR #48507:
URL: https://github.com/apache/airflow/pull/48507#discussion_r2025310246
##########
providers/databricks/src/airflow/providers/databricks/triggers/databricks.py:
##########
@@ -119,3 +119,84 @@ async def run(self):
}
)
return
+
+
+class DatabricksSQLStatementExecutionTrigger(BaseTrigger):
+ """
+ The trigger handles the logic of async communication with DataBricks SQL Statements API.
+
+ :param statement_id: ID of the SQL statement.
+ :param databricks_conn_id: Reference to the :ref:`Databricks connection <howto/connection:databricks>`.
+ :param polling_period_seconds: Controls the rate of the poll for the result of this run.
+ By default, the trigger will poll every 30 seconds.
+ :param retry_limit: The number of times to retry the connection in case of service outages.
+ :param retry_delay: The number of seconds to wait between retries.
+ :param retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` class.
+ """
+
+ def __init__(
+ self,
+ statement_id: str,
+ databricks_conn_id: str,
+ polling_period_seconds: int = 30,
+ retry_limit: int = 3,
+ retry_delay: int = 10,
+ retry_args: dict[Any, Any] | None = None,
+ caller: str = "DatabricksSQLStatementExecutionTrigger",
Review Comment:
I observed this caller is being set in all objects across the Databricks provider and needed for triggerer but more to know where the API request is coming from. It gets set in the hook initialisation if you look below and used in https://github.com/apache/airflow/blob/e85b61dd40ef6a1a16fd9260aff4fe6f228be6e8/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py#L157 which then gets set in the headers for each of API requests to Databricks as implemented there. So, I guess, it's good to keep it for consistency of the provider
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add DatabricksSQLStatementsOperator [airflow]
Posted by "pankajastro (via GitHub)" <gi...@apache.org>.
pankajastro commented on code in PR #48507:
URL: https://github.com/apache/airflow/pull/48507#discussion_r2024748340
##########
providers/databricks/src/airflow/providers/databricks/hooks/databricks.py:
##########
@@ -709,6 +771,47 @@ def update_job_permission(self, job_id: int, json: dict[str, Any]) -> dict:
"""
return self._do_api_call(("PATCH", f"api/2.0/permissions/jobs/{job_id}"), json)
+ def post_sql_statement(self, json: dict[str, Any]) -> str:
Review Comment:
just for consistency, shall we add docstring for param json?
##########
providers/databricks/src/airflow/providers/databricks/operators/databricks.py:
##########
@@ -969,6 +978,183 @@ def on_kill(self) -> None:
self.log.error("Error: Task: %s with invalid run_id was requested to be cancelled.", self.task_id)
+class DatabricksSQLStatementsOperator(BaseOperator):
+ """
+ Submits a Databricks SQL Statement to Databricks using the api/2.0/sql/statements/ API endpoint.
+
+ See: https://docs.databricks.com/api/workspace/statementexecution
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:DatabricksSQLStatementsOperator`
+
+ :param statement: The SQL statement to execute. The statement can optionally be parameterized, see parameters.
+ :param warehouse_id: Warehouse upon which to execute a statement.
+ :param catalog: Sets default catalog for statement execution, similar to USE CATALOG in SQL.
+ :param schema: Sets default schema for statement execution, similar to USE SCHEMA in SQL.
+ :param parameters: A list of parameters to pass into a SQL statement containing parameter markers.
+
+ .. seealso::
+ https://docs.databricks.com/api/workspace/statementexecution/executestatement#parameters
+ :param wait_for_termination: if we should wait for termination of the statement execution. ``True`` by default.
+ :param databricks_conn_id: Reference to the :ref:`Databricks connection <howto/connection:databricks>`.
+ By default and in the common case this will be ``databricks_default``. To use
+ token based authentication, provide the key ``token`` in the extra field for the
+ connection and create the key ``host`` and leave the ``host`` field empty. (templated)
+ :param polling_period_seconds: Controls the rate which we poll for the result of
+ this statement. By default the operator will poll every 30 seconds.
+ :param databricks_retry_limit: Amount of times retry if the Databricks backend is
+ unreachable. Its value must be greater than or equal to 1.
+ :param databricks_retry_delay: Number of seconds to wait between retries (it
+ might be a floating point number).
+ :param databricks_retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` class.
+ :param do_xcom_push: Whether we should push statement_id to xcom.
+ :param deferrable: Run operator in the deferrable mode.
+ """
+
+ # Used in airflow.models.BaseOperator
+ template_fields: Sequence[str] = ("databricks_conn_id",)
+ template_ext: Sequence[str] = (".json-tpl",)
+ # Databricks brand color (blue) under white text
+ ui_color = "#1CB1C2"
+ ui_fgcolor = "#fff"
+
+ def __init__(
+ self,
+ statement: str,
+ warehouse_id: str,
+ *,
+ catalog: str | None = None,
+ schema: str | None = None,
+ parameters: list[dict[str, Any]] | None = None,
+ databricks_conn_id: str = "databricks_default",
+ polling_period_seconds: int = 30,
+ databricks_retry_limit: int = 3,
+ databricks_retry_delay: int = 1,
+ databricks_retry_args: dict[Any, Any] | None = None,
+ do_xcom_push: bool = True,
+ wait_for_termination: bool = True,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ) -> None:
+ """Create a new ``DatabricksSubmitRunOperator``."""
+ super().__init__(**kwargs)
+ self.statement = statement
+ self.warehouse_id = warehouse_id
+ self.catalog = catalog
+ self.schema = schema
+ self.parameters = parameters
+ self.databricks_conn_id = databricks_conn_id
+ self.polling_period_seconds = polling_period_seconds
+ self.databricks_retry_limit = databricks_retry_limit
+ self.databricks_retry_delay = databricks_retry_delay
+ self.databricks_retry_args = databricks_retry_args
+ self.wait_for_termination = wait_for_termination
+ self.deferrable = deferrable
+
+ # This variable will be used in case our task gets killed.
+ self.statement_id: str | None = None
+ self.do_xcom_push = do_xcom_push
+
+ @cached_property
+ def _hook(self):
+ return self._get_hook(caller="DatabricksSQLStatementsOperator")
+
+ def _get_hook(self, caller: str) -> DatabricksHook:
+ return DatabricksHook(
+ self.databricks_conn_id,
+ retry_limit=self.databricks_retry_limit,
+ retry_delay=self.databricks_retry_delay,
+ retry_args=self.databricks_retry_args,
+ caller=caller,
+ )
+
+ def _handle_operator_execution(self) -> None:
+ while True:
+ statement_state = self._hook.get_sql_statement_state(self.statement_id)
+ if statement_state.is_terminal:
+ if statement_state.is_successful:
+ self.log.info("%s completed successfully.", self.task_id)
+ return
+ error_message = (
+ f"{self.task_id} failed with terminal state: {statement_state.state} "
+ f"and with the error code {statement_state.error_code} "
+ f"and error message {statement_state.error_message}"
+ )
+ raise AirflowException(error_message)
+
+ self.log.info("%s in run state: %s", self.task_id, statement_state.state)
+ self.log.info("Sleeping for %s seconds.", self.polling_period_seconds)
+ time.sleep(self.polling_period_seconds)
+
+ def _handle_deferrable_operator_execution(self) -> None:
+ statement_state = self._hook.get_sql_statement_state(self.statement_id)
+ if not statement_state.is_terminal:
+ self.defer(
+ trigger=DatabricksSQLStatementExecutionTrigger(
+ statement_id=self.statement_id,
+ databricks_conn_id=self.databricks_conn_id,
+ polling_period_seconds=self.polling_period_seconds,
+ retry_limit=self.databricks_retry_limit,
+ retry_delay=self.databricks_retry_delay,
+ retry_args=self.databricks_retry_args,
+ ),
+ method_name=DEFER_METHOD_NAME,
Review Comment:
what is `DEFER_METHOD_NAME`?
##########
providers/databricks/src/airflow/providers/databricks/operators/databricks.py:
##########
@@ -969,6 +978,183 @@ def on_kill(self) -> None:
self.log.error("Error: Task: %s with invalid run_id was requested to be cancelled.", self.task_id)
+class DatabricksSQLStatementsOperator(BaseOperator):
+ """
+ Submits a Databricks SQL Statement to Databricks using the api/2.0/sql/statements/ API endpoint.
+
+ See: https://docs.databricks.com/api/workspace/statementexecution
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:DatabricksSQLStatementsOperator`
+
+ :param statement: The SQL statement to execute. The statement can optionally be parameterized, see parameters.
+ :param warehouse_id: Warehouse upon which to execute a statement.
+ :param catalog: Sets default catalog for statement execution, similar to USE CATALOG in SQL.
+ :param schema: Sets default schema for statement execution, similar to USE SCHEMA in SQL.
+ :param parameters: A list of parameters to pass into a SQL statement containing parameter markers.
+
+ .. seealso::
+ https://docs.databricks.com/api/workspace/statementexecution/executestatement#parameters
+ :param wait_for_termination: if we should wait for termination of the statement execution. ``True`` by default.
+ :param databricks_conn_id: Reference to the :ref:`Databricks connection <howto/connection:databricks>`.
+ By default and in the common case this will be ``databricks_default``. To use
+ token based authentication, provide the key ``token`` in the extra field for the
+ connection and create the key ``host`` and leave the ``host`` field empty. (templated)
+ :param polling_period_seconds: Controls the rate which we poll for the result of
+ this statement. By default the operator will poll every 30 seconds.
+ :param databricks_retry_limit: Amount of times retry if the Databricks backend is
+ unreachable. Its value must be greater than or equal to 1.
+ :param databricks_retry_delay: Number of seconds to wait between retries (it
+ might be a floating point number).
+ :param databricks_retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` class.
+ :param do_xcom_push: Whether we should push statement_id to xcom.
+ :param deferrable: Run operator in the deferrable mode.
+ """
+
+ # Used in airflow.models.BaseOperator
+ template_fields: Sequence[str] = ("databricks_conn_id",)
+ template_ext: Sequence[str] = (".json-tpl",)
+ # Databricks brand color (blue) under white text
+ ui_color = "#1CB1C2"
+ ui_fgcolor = "#fff"
+
+ def __init__(
+ self,
+ statement: str,
+ warehouse_id: str,
+ *,
+ catalog: str | None = None,
+ schema: str | None = None,
+ parameters: list[dict[str, Any]] | None = None,
+ databricks_conn_id: str = "databricks_default",
+ polling_period_seconds: int = 30,
+ databricks_retry_limit: int = 3,
+ databricks_retry_delay: int = 1,
+ databricks_retry_args: dict[Any, Any] | None = None,
+ do_xcom_push: bool = True,
+ wait_for_termination: bool = True,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ) -> None:
+ """Create a new ``DatabricksSubmitRunOperator``."""
+ super().__init__(**kwargs)
+ self.statement = statement
+ self.warehouse_id = warehouse_id
+ self.catalog = catalog
+ self.schema = schema
+ self.parameters = parameters
+ self.databricks_conn_id = databricks_conn_id
+ self.polling_period_seconds = polling_period_seconds
+ self.databricks_retry_limit = databricks_retry_limit
+ self.databricks_retry_delay = databricks_retry_delay
+ self.databricks_retry_args = databricks_retry_args
+ self.wait_for_termination = wait_for_termination
+ self.deferrable = deferrable
+
+ # This variable will be used in case our task gets killed.
Review Comment:
These variables are getting used in execute method too so unsure if this comment makes sense or I'm missing anything?
##########
providers/databricks/src/airflow/providers/databricks/triggers/databricks.py:
##########
@@ -119,3 +119,84 @@ async def run(self):
}
)
return
+
+
+class DatabricksSQLStatementExecutionTrigger(BaseTrigger):
+ """
+ The trigger handles the logic of async communication with DataBricks SQL Statements API.
+
+ :param statement_id: ID of the SQL statement.
+ :param databricks_conn_id: Reference to the :ref:`Databricks connection <howto/connection:databricks>`.
+ :param polling_period_seconds: Controls the rate of the poll for the result of this run.
+ By default, the trigger will poll every 30 seconds.
+ :param retry_limit: The number of times to retry the connection in case of service outages.
+ :param retry_delay: The number of seconds to wait between retries.
+ :param retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` class.
+ """
+
+ def __init__(
+ self,
+ statement_id: str,
+ databricks_conn_id: str,
+ polling_period_seconds: int = 30,
+ retry_limit: int = 3,
+ retry_delay: int = 10,
+ retry_args: dict[Any, Any] | None = None,
+ caller: str = "DatabricksSQLStatementExecutionTrigger",
Review Comment:
Also, this param is not included in docstring
##########
providers/databricks/src/airflow/providers/databricks/triggers/databricks.py:
##########
@@ -119,3 +119,84 @@ async def run(self):
}
)
return
+
+
+class DatabricksSQLStatementExecutionTrigger(BaseTrigger):
+ """
+ The trigger handles the logic of async communication with DataBricks SQL Statements API.
+
+ :param statement_id: ID of the SQL statement.
+ :param databricks_conn_id: Reference to the :ref:`Databricks connection <howto/connection:databricks>`.
+ :param polling_period_seconds: Controls the rate of the poll for the result of this run.
+ By default, the trigger will poll every 30 seconds.
+ :param retry_limit: The number of times to retry the connection in case of service outages.
+ :param retry_delay: The number of seconds to wait between retries.
+ :param retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` class.
+ """
+
+ def __init__(
+ self,
+ statement_id: str,
+ databricks_conn_id: str,
+ polling_period_seconds: int = 30,
+ retry_limit: int = 3,
+ retry_delay: int = 10,
+ retry_args: dict[Any, Any] | None = None,
+ caller: str = "DatabricksSQLStatementExecutionTrigger",
Review Comment:
Shall we remove this?
I noticed that this parameter is not set during the initialization of DatabricksSQLStatementExecutionTrigger from the operator, nor is it included in the trigger's serialize method.
##########
providers/databricks/src/airflow/providers/databricks/operators/databricks.py:
##########
@@ -969,6 +978,183 @@ def on_kill(self) -> None:
self.log.error("Error: Task: %s with invalid run_id was requested to be cancelled.", self.task_id)
+class DatabricksSQLStatementsOperator(BaseOperator):
+ """
+ Submits a Databricks SQL Statement to Databricks using the api/2.0/sql/statements/ API endpoint.
+
+ See: https://docs.databricks.com/api/workspace/statementexecution
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:DatabricksSQLStatementsOperator`
+
+ :param statement: The SQL statement to execute. The statement can optionally be parameterized, see parameters.
+ :param warehouse_id: Warehouse upon which to execute a statement.
+ :param catalog: Sets default catalog for statement execution, similar to USE CATALOG in SQL.
+ :param schema: Sets default schema for statement execution, similar to USE SCHEMA in SQL.
+ :param parameters: A list of parameters to pass into a SQL statement containing parameter markers.
+
+ .. seealso::
+ https://docs.databricks.com/api/workspace/statementexecution/executestatement#parameters
+ :param wait_for_termination: if we should wait for termination of the statement execution. ``True`` by default.
+ :param databricks_conn_id: Reference to the :ref:`Databricks connection <howto/connection:databricks>`.
+ By default and in the common case this will be ``databricks_default``. To use
+ token based authentication, provide the key ``token`` in the extra field for the
+ connection and create the key ``host`` and leave the ``host`` field empty. (templated)
+ :param polling_period_seconds: Controls the rate which we poll for the result of
+ this statement. By default the operator will poll every 30 seconds.
+ :param databricks_retry_limit: Amount of times retry if the Databricks backend is
+ unreachable. Its value must be greater than or equal to 1.
+ :param databricks_retry_delay: Number of seconds to wait between retries (it
+ might be a floating point number).
+ :param databricks_retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` class.
+ :param do_xcom_push: Whether we should push statement_id to xcom.
+ :param deferrable: Run operator in the deferrable mode.
+ """
+
+ # Used in airflow.models.BaseOperator
+ template_fields: Sequence[str] = ("databricks_conn_id",)
+ template_ext: Sequence[str] = (".json-tpl",)
+ # Databricks brand color (blue) under white text
+ ui_color = "#1CB1C2"
+ ui_fgcolor = "#fff"
+
+ def __init__(
+ self,
+ statement: str,
+ warehouse_id: str,
+ *,
+ catalog: str | None = None,
+ schema: str | None = None,
+ parameters: list[dict[str, Any]] | None = None,
+ databricks_conn_id: str = "databricks_default",
+ polling_period_seconds: int = 30,
+ databricks_retry_limit: int = 3,
+ databricks_retry_delay: int = 1,
+ databricks_retry_args: dict[Any, Any] | None = None,
+ do_xcom_push: bool = True,
+ wait_for_termination: bool = True,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ) -> None:
+ """Create a new ``DatabricksSubmitRunOperator``."""
+ super().__init__(**kwargs)
+ self.statement = statement
+ self.warehouse_id = warehouse_id
+ self.catalog = catalog
+ self.schema = schema
+ self.parameters = parameters
+ self.databricks_conn_id = databricks_conn_id
+ self.polling_period_seconds = polling_period_seconds
+ self.databricks_retry_limit = databricks_retry_limit
+ self.databricks_retry_delay = databricks_retry_delay
+ self.databricks_retry_args = databricks_retry_args
+ self.wait_for_termination = wait_for_termination
+ self.deferrable = deferrable
+
+ # This variable will be used in case our task gets killed.
+ self.statement_id: str | None = None
+ self.do_xcom_push = do_xcom_push
+
+ @cached_property
+ def _hook(self):
+ return self._get_hook(caller="DatabricksSQLStatementsOperator")
+
+ def _get_hook(self, caller: str) -> DatabricksHook:
+ return DatabricksHook(
+ self.databricks_conn_id,
+ retry_limit=self.databricks_retry_limit,
+ retry_delay=self.databricks_retry_delay,
+ retry_args=self.databricks_retry_args,
+ caller=caller,
+ )
+
+ def _handle_operator_execution(self) -> None:
+ while True:
+ statement_state = self._hook.get_sql_statement_state(self.statement_id)
+ if statement_state.is_terminal:
+ if statement_state.is_successful:
+ self.log.info("%s completed successfully.", self.task_id)
+ return
+ error_message = (
+ f"{self.task_id} failed with terminal state: {statement_state.state} "
+ f"and with the error code {statement_state.error_code} "
+ f"and error message {statement_state.error_message}"
+ )
+ raise AirflowException(error_message)
+
+ self.log.info("%s in run state: %s", self.task_id, statement_state.state)
+ self.log.info("Sleeping for %s seconds.", self.polling_period_seconds)
+ time.sleep(self.polling_period_seconds)
+
+ def _handle_deferrable_operator_execution(self) -> None:
+ statement_state = self._hook.get_sql_statement_state(self.statement_id)
+ if not statement_state.is_terminal:
+ self.defer(
Review Comment:
shall we init timeout param too in self.defer
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add DatabricksSQLStatementsOperator [airflow]
Posted by "pankajkoti (via GitHub)" <gi...@apache.org>.
pankajkoti merged PR #48507:
URL: https://github.com/apache/airflow/pull/48507
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add DatabricksSQLStatementsOperator [airflow]
Posted by "rawwar (via GitHub)" <gi...@apache.org>.
rawwar commented on code in PR #48507:
URL: https://github.com/apache/airflow/pull/48507#discussion_r2024138039
##########
providers/databricks/src/airflow/providers/databricks/hooks/databricks.py:
##########
@@ -189,6 +189,65 @@ def from_json(cls, data: str) -> ClusterState:
return ClusterState(**json.loads(data))
+class SQLStatementState:
+ """Utility class for the SQL statement state concept of Databricks statements."""
+
+ SQL_STATEMENT_LIFE_CYCLE_STATES = [
+ "PENDING",
+ "RUNNING",
+ "SUCCEEDED",
+ "FAILED",
+ "CANCELED",
+ "CLOSED",
+ ]
+
+ def __init__(
+ self, state: str = "", error_code: str = "", error_message: str = "", *args, **kwargs
+ ) -> None:
+ if state not in self.SQL_STATEMENT_LIFE_CYCLE_STATES:
+ raise AirflowException(
+ f"Unexpected SQL statement life cycle state: {state}: If the state has "
+ "been introduced recently, please check the Databricks user "
+ "guide for troubleshooting information"
+ )
+
+ self.state = state
+ self.error_code = error_code
+ self.error_message = error_message
+
+ @property
+ def is_terminal(self) -> bool:
+ """True if the current state is a terminal state."""
+ return self.state in ("SUCCEEDED", "FAILED", "CANCELED", "CLOSED")
+
+ @property
+ def is_running(self) -> bool:
+ """True if the current state is running."""
+ return self.state in ("PENDING", "RUNNING")
+
+ @property
+ def is_successful(self) -> bool:
+ """True if the state is SUCCEEDED."""
+ return self.state == "SUCCEEDED"
+
+ def __eq__(self, other) -> bool:
+ return (
+ self.state == other.state
+ and self.error_code == other.error_code
+ and self.error_message == other.error_message
+ )
Review Comment:
```suggestion
def __eq__(self, other) -> bool:
if not isinstance(other, SQLStatementState):
return NotImplemented
return (
self.state == other.state
and self.error_code == other.error_code
and self.error_message == other.error_message
)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add DatabricksSQLStatementsOperator [airflow]
Posted by "pankajkoti (via GitHub)" <gi...@apache.org>.
pankajkoti commented on code in PR #48507:
URL: https://github.com/apache/airflow/pull/48507#discussion_r2025305634
##########
providers/databricks/src/airflow/providers/databricks/hooks/databricks.py:
##########
@@ -709,6 +771,47 @@ def update_job_permission(self, job_id: int, json: dict[str, Any]) -> dict:
"""
return self._do_api_call(("PATCH", f"api/2.0/permissions/jobs/{job_id}"), json)
+ def post_sql_statement(self, json: dict[str, Any]) -> str:
Review Comment:
added a docstring
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add DatabricksSQLStatementsOperator [airflow]
Posted by "pankajkoti (via GitHub)" <gi...@apache.org>.
pankajkoti commented on PR #48507:
URL: https://github.com/apache/airflow/pull/48507#issuecomment-2769785174
cc: @pankajastro @phanikumv @tatiana tagging for seeking review on the PR
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add DatabricksSQLStatementsOperator [airflow]
Posted by "rawwar (via GitHub)" <gi...@apache.org>.
rawwar commented on code in PR #48507:
URL: https://github.com/apache/airflow/pull/48507#discussion_r2024130960
##########
providers/databricks/src/airflow/providers/databricks/hooks/databricks.py:
##########
@@ -709,6 +768,47 @@ def update_job_permission(self, job_id: int, json: dict[str, Any]) -> dict:
"""
return self._do_api_call(("PATCH", f"api/2.0/permissions/jobs/{job_id}"), json)
+ def post_sql_statement(self, json: dict[str, Any]) -> str:
+ response = self._do_api_call(("POST", "api/2.0/sql/statements"), json)
+ return response["statement_id"]
+
+ def get_sql_statement_state(self, statement_id: str) -> SQLStatementState:
+ """
+ Retrieve run state of the SQL statement.
+
+ :param statement_id: ID of the SQL statement.
+ :return: state of the SQL statement.
+ """
+ get_statement_endpoint = ("GET", f"api/2.0/sql/statements/{statement_id}")
Review Comment:
At the top of this module, we are creating several constants for endpoint prefixes like `CANCEL_RUN_ENDPOINT` and etc. Maybe, we can create "STATEMENT_ENDPOINT" and reuse it across these new methods?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add DatabricksSQLStatementsOperator [airflow]
Posted by "phanikumv (via GitHub)" <gi...@apache.org>.
phanikumv commented on code in PR #48507:
URL: https://github.com/apache/airflow/pull/48507#discussion_r2024148643
##########
providers/databricks/docs/operators/sql_statements.rst:
##########
@@ -0,0 +1,57 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+.. _howto/operator:DatabricksSQLStatementsOperator:
+
+
+DatabricksSQLStatementsOperator
+===============================
+
+Use the :class:`~airflow.providers.databricks.operators.databricks.DatabricksSQLStatementsOperator` to submits a
Review Comment:
```suggestion
Use the :class:`~airflow.providers.databricks.operators.databricks.DatabricksSQLStatementsOperator` to submit a
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add DatabricksSQLStatementsOperator [airflow]
Posted by "pankajkoti (via GitHub)" <gi...@apache.org>.
pankajkoti commented on code in PR #48507:
URL: https://github.com/apache/airflow/pull/48507#discussion_r2025313956
##########
providers/databricks/src/airflow/providers/databricks/operators/databricks.py:
##########
@@ -969,6 +978,183 @@ def on_kill(self) -> None:
self.log.error("Error: Task: %s with invalid run_id was requested to be cancelled.", self.task_id)
+class DatabricksSQLStatementsOperator(BaseOperator):
+ """
+ Submits a Databricks SQL Statement to Databricks using the api/2.0/sql/statements/ API endpoint.
+
+ See: https://docs.databricks.com/api/workspace/statementexecution
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:DatabricksSQLStatementsOperator`
+
+ :param statement: The SQL statement to execute. The statement can optionally be parameterized, see parameters.
+ :param warehouse_id: Warehouse upon which to execute a statement.
+ :param catalog: Sets default catalog for statement execution, similar to USE CATALOG in SQL.
+ :param schema: Sets default schema for statement execution, similar to USE SCHEMA in SQL.
+ :param parameters: A list of parameters to pass into a SQL statement containing parameter markers.
+
+ .. seealso::
+ https://docs.databricks.com/api/workspace/statementexecution/executestatement#parameters
+ :param wait_for_termination: if we should wait for termination of the statement execution. ``True`` by default.
+ :param databricks_conn_id: Reference to the :ref:`Databricks connection <howto/connection:databricks>`.
+ By default and in the common case this will be ``databricks_default``. To use
+ token based authentication, provide the key ``token`` in the extra field for the
+ connection and create the key ``host`` and leave the ``host`` field empty. (templated)
+ :param polling_period_seconds: Controls the rate which we poll for the result of
+ this statement. By default the operator will poll every 30 seconds.
+ :param databricks_retry_limit: Amount of times retry if the Databricks backend is
+ unreachable. Its value must be greater than or equal to 1.
+ :param databricks_retry_delay: Number of seconds to wait between retries (it
+ might be a floating point number).
+ :param databricks_retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` class.
+ :param do_xcom_push: Whether we should push statement_id to xcom.
+ :param deferrable: Run operator in the deferrable mode.
+ """
+
+ # Used in airflow.models.BaseOperator
+ template_fields: Sequence[str] = ("databricks_conn_id",)
+ template_ext: Sequence[str] = (".json-tpl",)
+ # Databricks brand color (blue) under white text
+ ui_color = "#1CB1C2"
+ ui_fgcolor = "#fff"
+
+ def __init__(
+ self,
+ statement: str,
+ warehouse_id: str,
+ *,
+ catalog: str | None = None,
+ schema: str | None = None,
+ parameters: list[dict[str, Any]] | None = None,
+ databricks_conn_id: str = "databricks_default",
+ polling_period_seconds: int = 30,
+ databricks_retry_limit: int = 3,
+ databricks_retry_delay: int = 1,
+ databricks_retry_args: dict[Any, Any] | None = None,
+ do_xcom_push: bool = True,
+ wait_for_termination: bool = True,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ) -> None:
+ """Create a new ``DatabricksSubmitRunOperator``."""
+ super().__init__(**kwargs)
+ self.statement = statement
+ self.warehouse_id = warehouse_id
+ self.catalog = catalog
+ self.schema = schema
+ self.parameters = parameters
+ self.databricks_conn_id = databricks_conn_id
+ self.polling_period_seconds = polling_period_seconds
+ self.databricks_retry_limit = databricks_retry_limit
+ self.databricks_retry_delay = databricks_retry_delay
+ self.databricks_retry_args = databricks_retry_args
+ self.wait_for_termination = wait_for_termination
+ self.deferrable = deferrable
+
+ # This variable will be used in case our task gets killed.
+ self.statement_id: str | None = None
+ self.do_xcom_push = do_xcom_push
+
+ @cached_property
+ def _hook(self):
+ return self._get_hook(caller="DatabricksSQLStatementsOperator")
+
+ def _get_hook(self, caller: str) -> DatabricksHook:
+ return DatabricksHook(
+ self.databricks_conn_id,
+ retry_limit=self.databricks_retry_limit,
+ retry_delay=self.databricks_retry_delay,
+ retry_args=self.databricks_retry_args,
+ caller=caller,
+ )
+
+ def _handle_operator_execution(self) -> None:
+ while True:
+ statement_state = self._hook.get_sql_statement_state(self.statement_id)
+ if statement_state.is_terminal:
+ if statement_state.is_successful:
+ self.log.info("%s completed successfully.", self.task_id)
+ return
+ error_message = (
+ f"{self.task_id} failed with terminal state: {statement_state.state} "
+ f"and with the error code {statement_state.error_code} "
+ f"and error message {statement_state.error_message}"
+ )
+ raise AirflowException(error_message)
+
+ self.log.info("%s in run state: %s", self.task_id, statement_state.state)
+ self.log.info("Sleeping for %s seconds.", self.polling_period_seconds)
+ time.sleep(self.polling_period_seconds)
+
+ def _handle_deferrable_operator_execution(self) -> None:
+ statement_state = self._hook.get_sql_statement_state(self.statement_id)
+ if not statement_state.is_terminal:
+ self.defer(
Review Comment:
Good idea! I noticed that most providers including Databricks and operators currently lack this feature. However, I’ve now introduced a timeout for both synchronous and asynchronous execution in this newly added operator.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add DatabricksSQLStatementsOperator [airflow]
Posted by "pankajkoti (via GitHub)" <gi...@apache.org>.
pankajkoti commented on code in PR #48507:
URL: https://github.com/apache/airflow/pull/48507#discussion_r2025316015
##########
providers/databricks/src/airflow/providers/databricks/operators/databricks.py:
##########
@@ -969,6 +978,183 @@ def on_kill(self) -> None:
self.log.error("Error: Task: %s with invalid run_id was requested to be cancelled.", self.task_id)
+class DatabricksSQLStatementsOperator(BaseOperator):
+ """
+ Submits a Databricks SQL Statement to Databricks using the api/2.0/sql/statements/ API endpoint.
+
+ See: https://docs.databricks.com/api/workspace/statementexecution
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:DatabricksSQLStatementsOperator`
+
+ :param statement: The SQL statement to execute. The statement can optionally be parameterized, see parameters.
+ :param warehouse_id: Warehouse upon which to execute a statement.
+ :param catalog: Sets default catalog for statement execution, similar to USE CATALOG in SQL.
+ :param schema: Sets default schema for statement execution, similar to USE SCHEMA in SQL.
+ :param parameters: A list of parameters to pass into a SQL statement containing parameter markers.
+
+ .. seealso::
+ https://docs.databricks.com/api/workspace/statementexecution/executestatement#parameters
+ :param wait_for_termination: if we should wait for termination of the statement execution. ``True`` by default.
+ :param databricks_conn_id: Reference to the :ref:`Databricks connection <howto/connection:databricks>`.
+ By default and in the common case this will be ``databricks_default``. To use
+ token based authentication, provide the key ``token`` in the extra field for the
+ connection and create the key ``host`` and leave the ``host`` field empty. (templated)
+ :param polling_period_seconds: Controls the rate which we poll for the result of
+ this statement. By default the operator will poll every 30 seconds.
+ :param databricks_retry_limit: Amount of times retry if the Databricks backend is
+ unreachable. Its value must be greater than or equal to 1.
+ :param databricks_retry_delay: Number of seconds to wait between retries (it
+ might be a floating point number).
+ :param databricks_retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` class.
+ :param do_xcom_push: Whether we should push statement_id to xcom.
+ :param deferrable: Run operator in the deferrable mode.
+ """
+
+ # Used in airflow.models.BaseOperator
+ template_fields: Sequence[str] = ("databricks_conn_id",)
+ template_ext: Sequence[str] = (".json-tpl",)
+ # Databricks brand color (blue) under white text
+ ui_color = "#1CB1C2"
+ ui_fgcolor = "#fff"
+
+ def __init__(
+ self,
+ statement: str,
+ warehouse_id: str,
+ *,
+ catalog: str | None = None,
+ schema: str | None = None,
+ parameters: list[dict[str, Any]] | None = None,
+ databricks_conn_id: str = "databricks_default",
+ polling_period_seconds: int = 30,
+ databricks_retry_limit: int = 3,
+ databricks_retry_delay: int = 1,
+ databricks_retry_args: dict[Any, Any] | None = None,
+ do_xcom_push: bool = True,
+ wait_for_termination: bool = True,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ) -> None:
+ """Create a new ``DatabricksSubmitRunOperator``."""
+ super().__init__(**kwargs)
+ self.statement = statement
+ self.warehouse_id = warehouse_id
+ self.catalog = catalog
+ self.schema = schema
+ self.parameters = parameters
+ self.databricks_conn_id = databricks_conn_id
+ self.polling_period_seconds = polling_period_seconds
+ self.databricks_retry_limit = databricks_retry_limit
+ self.databricks_retry_delay = databricks_retry_delay
+ self.databricks_retry_args = databricks_retry_args
+ self.wait_for_termination = wait_for_termination
+ self.deferrable = deferrable
+
+ # This variable will be used in case our task gets killed.
Review Comment:
The comment does not intend to talk about all "these" variables, but just "this" particular variable self.statement_id. To make it clearer, I added another blank line after the variable(we already have had a blank line separator before this variable & the comment) to further differentiate.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add DatabricksSQLStatementsOperator [airflow]
Posted by "pankajkoti (via GitHub)" <gi...@apache.org>.
pankajkoti commented on code in PR #48507:
URL: https://github.com/apache/airflow/pull/48507#discussion_r2024273364
##########
providers/databricks/src/airflow/providers/databricks/hooks/databricks.py:
##########
@@ -189,6 +189,65 @@ def from_json(cls, data: str) -> ClusterState:
return ClusterState(**json.loads(data))
+class SQLStatementState:
+ """Utility class for the SQL statement state concept of Databricks statements."""
+
+ SQL_STATEMENT_LIFE_CYCLE_STATES = [
+ "PENDING",
+ "RUNNING",
+ "SUCCEEDED",
+ "FAILED",
+ "CANCELED",
+ "CLOSED",
+ ]
+
+ def __init__(
+ self, state: str = "", error_code: str = "", error_message: str = "", *args, **kwargs
+ ) -> None:
+ if state not in self.SQL_STATEMENT_LIFE_CYCLE_STATES:
+ raise AirflowException(
+ f"Unexpected SQL statement life cycle state: {state}: If the state has "
+ "been introduced recently, please check the Databricks user "
+ "guide for troubleshooting information"
+ )
+
+ self.state = state
+ self.error_code = error_code
+ self.error_message = error_message
+
+ @property
+ def is_terminal(self) -> bool:
+ """True if the current state is a terminal state."""
+ return self.state in ("SUCCEEDED", "FAILED", "CANCELED", "CLOSED")
+
+ @property
+ def is_running(self) -> bool:
+ """True if the current state is running."""
+ return self.state in ("PENDING", "RUNNING")
+
+ @property
+ def is_successful(self) -> bool:
+ """True if the state is SUCCEEDED."""
+ return self.state == "SUCCEEDED"
+
+ def __eq__(self, other) -> bool:
+ return (
+ self.state == other.state
+ and self.error_code == other.error_code
+ and self.error_message == other.error_message
+ )
Review Comment:
added similar to RunState
##########
providers/databricks/docs/operators/sql_statements.rst:
##########
@@ -0,0 +1,57 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+.. _howto/operator:DatabricksSQLStatementsOperator:
+
+
+DatabricksSQLStatementsOperator
+===============================
+
+Use the :class:`~airflow.providers.databricks.operators.databricks.DatabricksSQLStatementsOperator` to submits a
Review Comment:
corrected
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add DatabricksSQLStatementsOperator [airflow]
Posted by "pankajkoti (via GitHub)" <gi...@apache.org>.
pankajkoti commented on code in PR #48507:
URL: https://github.com/apache/airflow/pull/48507#discussion_r2024274065
##########
providers/databricks/src/airflow/providers/databricks/hooks/databricks.py:
##########
@@ -709,6 +768,47 @@ def update_job_permission(self, job_id: int, json: dict[str, Any]) -> dict:
"""
return self._do_api_call(("PATCH", f"api/2.0/permissions/jobs/{job_id}"), json)
+ def post_sql_statement(self, json: dict[str, Any]) -> str:
+ response = self._do_api_call(("POST", "api/2.0/sql/statements"), json)
+ return response["statement_id"]
+
+ def get_sql_statement_state(self, statement_id: str) -> SQLStatementState:
+ """
+ Retrieve run state of the SQL statement.
+
+ :param statement_id: ID of the SQL statement.
+ :return: state of the SQL statement.
+ """
+ get_statement_endpoint = ("GET", f"api/2.0/sql/statements/{statement_id}")
Review Comment:
added such a constant and used it
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org