Source code for airflow.providers.apache.spark.hooks.spark_sql

#
# 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.
from __future__ import annotations

import subprocess
from typing import TYPE_CHECKING, Any

from airflow.exceptions import AirflowException, AirflowNotFoundException
from airflow.hooks.base import BaseHook

if TYPE_CHECKING:
    from airflow.models.connection import Connection


[docs]class SparkSqlHook(BaseHook): """ This hook is a wrapper around the spark-sql binary; requires the "spark-sql" binary to be in the PATH. :param sql: The SQL query to execute :param conf: arbitrary Spark configuration property :param conn_id: connection_id string :param total_executor_cores: (Standalone & Mesos only) Total cores for all executors (Default: all the available cores on the worker) :param executor_cores: (Standalone & YARN only) Number of cores per executor (Default: 2) :param executor_memory: Memory per executor (e.g. 1000M, 2G) (Default: 1G) :param keytab: Full path to the file that contains the keytab :param master: spark://host:port, mesos://host:port, yarn, or local (Default: The ``host`` and ``port`` set in the Connection, or ``"yarn"``) :param name: Name of the job. :param num_executors: Number of executors to launch :param verbose: Whether to pass the verbose flag to spark-sql :param yarn_queue: The YARN queue to submit to (Default: The ``queue`` value set in the Connection, or ``"default"``) """
[docs] conn_name_attr = "conn_id"
[docs] default_conn_name = "spark_sql_default"
[docs] conn_type = "spark_sql"
[docs] hook_name = "Spark SQL"
@classmethod
[docs] def get_ui_field_behaviour(cls) -> dict[str, Any]: """Return custom UI field behaviour for Spark SQL connection.""" return { "hidden_fields": ["schema", "login", "password", "extra"], "relabeling": {}, }
@classmethod
[docs] def get_connection_form_widgets(cls) -> dict[str, Any]: """Return connection widgets to add to Spark SQL connection form.""" from flask_appbuilder.fieldwidgets import BS3TextFieldWidget from flask_babel import lazy_gettext from wtforms import StringField from wtforms.validators import Optional return { "queue": StringField( lazy_gettext("YARN queue"), widget=BS3TextFieldWidget(), description="Default YARN queue to use", validators=[Optional()], ) }
def __init__( self, sql: str, conf: str | None = None, conn_id: str = default_conn_name, total_executor_cores: int | None = None, executor_cores: int | None = None, executor_memory: str | None = None, keytab: str | None = None, principal: str | None = None, master: str | None = None, name: str = "default-name", num_executors: int | None = None, verbose: bool = True, yarn_queue: str | None = None, ) -> None: super().__init__() options: dict = {} conn: Connection | None = None try: conn = self.get_connection(conn_id) except AirflowNotFoundException: conn = None if conn: options = conn.extra_dejson # Set arguments to values set in Connection if not explicitly provided. if master is None: if conn is None: master = "yarn" elif conn.port: master = f"{conn.host}:{conn.port}" else: master = conn.host if yarn_queue is None: yarn_queue = options.get("queue", "default") self._sql = sql self._conf = conf self._total_executor_cores = total_executor_cores self._executor_cores = executor_cores self._executor_memory = executor_memory self._keytab = keytab self._principal = principal self._master = master self._name = name self._num_executors = num_executors self._verbose = verbose self._yarn_queue = yarn_queue self._sp: Any = None
[docs] def get_conn(self) -> Any: pass
def _prepare_command(self, cmd: str | list[str]) -> list[str]: """ Construct the spark-sql command to execute. Verbose output is enabled as default. :param cmd: command to append to the spark-sql command :return: full command to be executed """ connection_cmd = ["spark-sql"] if self._conf: for conf_el in self._conf.split(","): connection_cmd += ["--conf", conf_el] if self._total_executor_cores: connection_cmd += ["--total-executor-cores", str(self._total_executor_cores)] if self._executor_cores: connection_cmd += ["--executor-cores", str(self._executor_cores)] if self._executor_memory: connection_cmd += ["--executor-memory", self._executor_memory] if self._keytab: connection_cmd += ["--keytab", self._keytab] if self._principal: connection_cmd += ["--principal", self._principal] if self._num_executors: connection_cmd += ["--num-executors", str(self._num_executors)] if self._sql: sql = self._sql.strip() if sql.endswith((".sql", ".hql")): connection_cmd += ["-f", sql] else: connection_cmd += ["-e", sql] if self._master: connection_cmd += ["--master", self._master] if self._name: connection_cmd += ["--name", self._name] if self._verbose: connection_cmd += ["--verbose"] if self._yarn_queue: connection_cmd += ["--queue", self._yarn_queue] if isinstance(cmd, str): connection_cmd += cmd.split() elif isinstance(cmd, list): connection_cmd += cmd else: raise AirflowException(f"Invalid additional command: {cmd}") self.log.debug("Spark-Sql cmd: %s", connection_cmd) return connection_cmd
[docs] def run_query(self, cmd: str = "", **kwargs: Any) -> None: """ Remote Popen (actually execute the Spark-sql query). :param cmd: command to append to the spark-sql command :param kwargs: extra arguments to Popen (see subprocess.Popen) """ spark_sql_cmd = self._prepare_command(cmd) self._sp = subprocess.Popen( spark_sql_cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True, **kwargs ) for line in iter(self._sp.stdout): # type: ignore self.log.info(line) returncode = self._sp.wait() if returncode: raise AirflowException( f"Cannot execute '{self._sql}' on {self._master} (additional parameters: '{cmd}'). " f"Process exit code: {returncode}." )
[docs] def kill(self) -> None: """Kill Spark job.""" if self._sp and self._sp.poll() is None: self.log.info("Killing the Spark-Sql job") self._sp.kill()

Was this entry helpful?