[AIRFLOW-3888] HA for metastore connection (apache#4708)

GuzikJakub · Fokko · commit 355bd56282e6 · 2019-05-10T12:45:55.000+02:00
* HA for Metastore

* [AIRFLOW-3888] HA for metastore connection

Creating a connection to a metasotor with two hosts for high avitablity (eg connection 1, connection 2) is not possible because the entire value entered is taken. For our needs, it is necessary to go through subsequent hosts and connect to the first working.

This change allows you to check and then connect to a working metastor.

* add function to base_hook

* update webhdfs_hook

* back to original version

* back to original version

* Update hive_hooks.py

Thank you. I made a few changes because during the tests I detected several errors.

I have a question, when I do marge to my pull it will be  still possible to land it in the airflow main branch?

* [AIRFLOW-3888] HA for metastore connection 

flake8 code repair

* [AIRFLOW-3888] HA for metastore connection 

Flake8 repair

* [AIRFLOW-3888] HA for metastore connection

Code behavior improvements

* [AIRFLOW-3888] HA for metastore connection 

Add test

* [AIRFLOW-3888] HA for metastore connection

test improvement

* [AIRFLOW-3888] HA for metastore connection

Add test

[AIRFLOW-3888] HA for metastore connection

test improvement

* [AIRFLOW-3888] HA for metastore connection

Add test

[AIRFLOW-3888] HA for metastore connection

test improvement

[AIRFLOW-3888] HA for metastore connection

test improvement

* [AIRFLOW-3888] HA for metastore connection

Improving the typo in the variable name

*  [AIRFLOW-3888] HA for metastore connection

Mock return_value edit

* [AIRFLOW-3888] HA for metastore connection

Flake8 repair

* [AIRFLOW-3888] HA for metastore connection

Test repair

* [AIRFLOW-3888] HA for metastore connection

Flake8 repair

[AIRFLOW-3888] HA for metastore connection

Test repair
diff --git a/airflow/hooks/hive_hooks.py b/airflow/hooks/hive_hooks.py
@@ -22,6 +22,7 @@
 import re
 import subprocess
 import time
+import socket
 from collections import OrderedDict
 from tempfile import NamedTemporaryFile
 
@@ -471,7 +472,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):
@@ -492,13 +493,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:
@@ -514,14 +522,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
@@ -33,6 +33,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
@@ -283,6 +284,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)
 
