[airflow] branch v1-10-stable updated: Functionality to shuffle HMS connections used by HiveMetastoreHook facilitating load balancing (#9280)

[ka...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

kaxilnaik pushed a commit to branch v1-10-stable
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/v1-10-stable by this push:
     new 868945d  Functionality to shuffle HMS connections used by HiveMetastoreHook facilitating load balancing (#9280)
868945d is described below

commit 868945d6ec2a2629f73f6f5e798a139e951908e0
Author: vanka56 <ve...@gmail.com>
AuthorDate: Tue Jul 7 03:11:40 2020 -0700

    Functionality to shuffle HMS connections used by HiveMetastoreHook facilitating load balancing (#9280)
    
    Co-authored-by: Venkatesh Selvaraj <ve...@pinterest.com>
---
 airflow/hooks/hive_hooks.py   |  2 ++
 tests/hooks/test_hive_hook.py | 11 ++++++++---
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/airflow/hooks/hive_hooks.py b/airflow/hooks/hive_hooks.py
index ec37dfb..e521d7b 100644
--- a/airflow/hooks/hive_hooks.py
+++ b/airflow/hooks/hive_hooks.py
@@ -22,6 +22,7 @@ from __future__ import print_function, unicode_literals
 import collections
 import contextlib
 import os
+import random
 import re
 import subprocess
 import time
@@ -550,6 +551,7 @@ class HiveMetastoreHook(BaseHook):
 
     def _find_valid_server(self):
         conns = self.get_connections(self.conn_id)
+        random.shuffle(conns)
         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)
diff --git a/tests/hooks/test_hive_hook.py b/tests/hooks/test_hive_hook.py
index 011038a..025d421 100644
--- a/tests/hooks/test_hive_hook.py
+++ b/tests/hooks/test_hive_hook.py
@@ -35,7 +35,7 @@ from airflow.models.dag import DAG
 from airflow.secrets.environment_variables import CONN_ENV_PREFIX
 from airflow.utils import timezone
 from airflow.utils.operator_helpers import AIRFLOW_VAR_NAME_FORMAT_MAPPING
-from tests.compat import patch, MagicMock, call
+from tests.compat import patch, MagicMock, call, mock
 from tests.test_utils.mock_hooks import MockHiveCliHook, MockHiveServer2Hook
 from tests.test_utils.mock_process import MockSubProcess
 
@@ -553,8 +553,13 @@ class TestHiveMetastoreHook(TestHiveEnvironment):
         self.assertFalse(
             self.hook.table_exists("does-not-exist")
         )
-        self.hook.metastore.__enter__().get_table.assert_called_with(
-            dbname='default', tbl_name='does-not-exist')
+
+    @mock.patch('airflow.hooks.hive_hooks.random.shuffle')
+    @mock.patch('airflow.hooks.hive_hooks.HiveMetastoreHook.get_connections')
+    @mock.patch('airflow.hooks.hive_hooks.HiveMetastoreHook.get_metastore_client')
+    def test_check_hms_clients_load_balance(self, mock_client, mock_get_conn, mock_shuffle):
+        HiveMetastoreHook()._find_valid_server()
+        mock_shuffle.assert_called_once_with(mock_get_conn.return_value)
 
 
 class TestHiveServer2Hook(unittest.TestCase):