diff --git a/airflow/hooks/hive_hooks.py b/airflow/hooks/hive_hooks.py index 2efeb3408a9c9..e718dd4096a8a 100644 --- a/airflow/hooks/hive_hooks.py +++ b/airflow/hooks/hive_hooks.py @@ -24,6 +24,7 @@ import re import subprocess import time +import socket from collections import OrderedDict from tempfile import NamedTemporaryFile @@ -477,7 +478,7 @@ class HiveMetastoreHook(BaseHook): MAX_PART_COUNT = 32767 def __init__(self, metastore_conn_id='metastore_default'): - self.metastore_conn = self.get_connection(metastore_conn_id) + self.conn_id = metastore_conn_id self.metastore = self.get_metastore_client() def __getstate__(self): @@ -498,13 +499,20 @@ def get_metastore_client(self): import hmsclient from thrift.transport import TSocket, TTransport from thrift.protocol import TBinaryProtocol - ms = self.metastore_conn + + ms = self._find_valid_server() + + if ms is None: + raise AirflowException("Failed to locate the valid server.") + auth_mechanism = ms.extra_dejson.get('authMechanism', 'NOSASL') + if configuration.conf.get('core', 'security') == 'kerberos': auth_mechanism = ms.extra_dejson.get('authMechanism', 'GSSAPI') kerberos_service_name = ms.extra_dejson.get('kerberos_service_name', 'hive') - socket = TSocket.TSocket(ms.host, ms.port) + conn_socket = TSocket.TSocket(ms.host, ms.port) + if configuration.conf.get('core', 'security') == 'kerberos' \ and auth_mechanism == 'GSSAPI': try: @@ -520,14 +528,26 @@ def sasl_factory(): return sasl_client from thrift_sasl import TSaslClientTransport - transport = TSaslClientTransport(sasl_factory, "GSSAPI", socket) + transport = TSaslClientTransport(sasl_factory, "GSSAPI", conn_socket) else: - transport = TTransport.TBufferedTransport(socket) + transport = TTransport.TBufferedTransport(conn_socket) protocol = TBinaryProtocol.TBinaryProtocol(transport) return hmsclient.HMSClient(iprot=protocol) + def _find_valid_server(self): + conns = self.get_connections(self.conn_id) + for conn in conns: + host_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self.log.info("Trying to connect to %s:%s", conn.host, conn.port) + if host_socket.connect_ex((conn.host, conn.port)) == 0: + self.log.info("Connected to %s:%s", conn.host, conn.port) + host_socket.close() + return conn + else: + self.log.info("Could not connect to %s:%s", conn.host, conn.port) + def get_conn(self): return self.metastore diff --git a/tests/hooks/test_hive_hook.py b/tests/hooks/test_hive_hook.py index 24bed195994d1..888de8f2bb731 100644 --- a/tests/hooks/test_hive_hook.py +++ b/tests/hooks/test_hive_hook.py @@ -32,6 +32,7 @@ from airflow import DAG, configuration from airflow.exceptions import AirflowException from airflow.hooks.hive_hooks import HiveCliHook, HiveMetastoreHook, HiveServer2Hook +from airflow.models.connection import Connection from airflow.operators.hive_operator import HiveOperator from airflow.utils import timezone from airflow.utils.operator_helpers import AIRFLOW_VAR_NAME_FORMAT_MAPPING @@ -282,6 +283,13 @@ def test_get_max_partition_from_valid_part_specs(self): def test_get_metastore_client(self): self.assertIsInstance(self.hook.get_metastore_client(), HMSClient) + @mock.patch("airflow.hooks.hive_hooks.HiveMetastoreHook.get_connection", + return_value=[Connection(host="localhost", port="9802")]) + @mock.patch("airflow.hooks.hive_hooks.socket") + def test_error_metastore_client(self, socket_mock, _find_vaild_server_mock): + socket_mock.socket.return_value.connect_ex.return_value = 0 + self.hook.get_metastore_client() + def test_get_conn(self): self.assertIsInstance(self.hook.get_conn(), HMSClient)