浏览代码

HUE-6160 [libsentry] Cache Sentry client and update tests to include api2

Jenny Kim 8 年之前
父节点
当前提交
81e6db4

+ 25 - 4
desktop/libs/libsentry/src/libsentry/api.py

@@ -16,6 +16,7 @@
 # limitations under the License.
 
 import logging
+import threading
 import time
 
 from django.utils.translation import ugettext as _
@@ -29,6 +30,9 @@ from libsentry.sentry_site import get_sentry_client, is_ha_enabled
 
 LOG = logging.getLogger(__name__)
 
+API_CACHE = None
+API_CACHE_LOCK = threading.Lock()
+
 
 def ha_error_handler(func):
   def decorator(*args, **kwargs):
@@ -44,8 +48,7 @@ def ha_error_handler(func):
         else:
           LOG.info('Could not connect to Sentry server %s, attempting to fetch next available client.' % args[0].client.host)
           time.sleep(1)
-          args[0].client = get_sentry_client(args[0].client.username, SentryClient)
-          LOG.info('Picked %s' % args[0].client)
+          args[0].client = get_cached_client(args[0].client.username, force_reset=True)
       except SentryException, e:
         raise e
       except Exception, e:
@@ -55,11 +58,29 @@ def ha_error_handler(func):
 
 
 def get_api(user):
-  client = get_sentry_client(user.username, SentryClient)
-
+  client = get_cached_client(user.username)
   return SentryApi(client)
 
 
+def get_cached_client(username, force_reset=False):
+  exempt_host = None
+
+  global API_CACHE
+  if force_reset and API_CACHE is not None:
+    exempt_host = API_CACHE.host
+    LOG.info("Force resetting the cached Sentry client to exempt current host: %s" % exempt_host)
+    API_CACHE = None
+
+  if API_CACHE is None:
+    API_CACHE_LOCK.acquire()
+    try:
+      API_CACHE = get_sentry_client(username, SentryClient, exempt_host=exempt_host)
+      LOG.info("Setting cached Sentry client to host: %s" % API_CACHE.host)
+    finally:
+      API_CACHE_LOCK.release()
+  return API_CACHE
+
+
 class SentryApi(object):
 
   def __init__(self, client):

+ 26 - 2
desktop/libs/libsentry/src/libsentry/api2.py

@@ -16,6 +16,7 @@
 # limitations under the License.
 
 import logging
+import threading
 import time
 
 from django.utils.translation import ugettext as _
@@ -29,6 +30,9 @@ from libsentry.sentry_site import get_sentry_client, is_ha_enabled
 
 LOG = logging.getLogger(__name__)
 
+API_CACHE = None
+API_CACHE_LOCK = threading.Lock()
+
 
 def ha_error_handler(func):
   def decorator(*args, **kwargs):
@@ -44,7 +48,7 @@ def ha_error_handler(func):
         else:
           LOG.info('Could not connect to Sentry server %s, attempting to fetch next available client.' % args[0].client.host)
           time.sleep(1)
-          args[0].client = get_sentry_client(args[0].client.username, SentryClient)
+          args[0].client = get_cached_client(args[0].client.username, force_reset=True)
           LOG.info('Picked %s' % args[0].client)
       except SentryException, e:
         raise e
@@ -58,11 +62,31 @@ def get_api(user, component):
   if component == 'solr':
     component = component.upper()
 
-  client = get_sentry_client(user.username, SentryClient, component)
+  client = get_cached_client(user.username, component=component)
 
   return SentryApi(client)
 
 
+def get_cached_client(username, component=None, force_reset=False):
+  exempt_host = None
+
+  global API_CACHE
+  if force_reset and API_CACHE is not None:
+    exempt_host = API_CACHE.host
+    component = API_CACHE.component
+    LOG.info("Force resetting the cached Sentry client for component %s to exempt current host: %s" % (component, exempt_host))
+    API_CACHE = None
+
+  if API_CACHE is None:
+    API_CACHE_LOCK.acquire()
+    try:
+      API_CACHE = get_sentry_client(username, SentryClient, exempt_host=exempt_host, component=component)
+      LOG.info("Setting cached Sentry client to host: %s" % API_CACHE.host)
+    finally:
+      API_CACHE_LOCK.release()
+  return API_CACHE
+
+
 class SentryApi(object):
 
   def __init__(self, client):

+ 5 - 4
desktop/libs/libsentry/src/libsentry/sentry_site.py

@@ -98,10 +98,10 @@ def is_ha_enabled():
   return get_sentry_server_rpc_addresses() is not None
 
 
-def get_sentry_client(username, client_class, component=None):
+def get_sentry_client(username, client_class, exempt_host=None, component=None):
   server = None
   if is_ha_enabled():
-    servers = _get_server_properties()
+    servers = _get_server_properties(exempt_host=exempt_host)
     if servers:
       server = random.choice(servers)
 
@@ -124,12 +124,13 @@ def get_sentry_client(username, client_class, component=None):
   return client
 
 
-def _get_server_properties():
+def _get_server_properties(exempt_host=None):
   try:
     servers = []
     sentry_servers = get_sentry_server_rpc_addresses()
     for server in sentry_servers:
-      servers.append({'hostname': server, 'port': get_sentry_server_rpc_port()})
+      if server != exempt_host:
+        servers.append({'hostname': server, 'port': get_sentry_server_rpc_port()})
   except Exception, e:
     raise PopupException(_('Error in retrieving Sentry server properties.'), detail=e)
 

+ 99 - 85
desktop/libs/libsentry/src/libsentry/tests.py

@@ -32,20 +32,18 @@ from desktop.lib.test_utils import add_to_group, grant_access
 from hadoop.pseudo_hdfs4 import is_live_cluster
 
 from libsentry import sentry_site
-from libsentry.api import get_api
-from libsentry.api2 import get_api as get_api2
+from libsentry.api import get_api, API_CACHE
+from libsentry.api2 import get_api as get_api2, API_CACHE as API2_CACHE
 from libsentry.conf import is_enabled, HOSTNAME, PORT, SENTRY_CONF_DIR
 from libsentry.client import SentryClient
 
 
-class TestWithSentry:
+class TestWithSentry(object):
 
   requires_hadoop = True
 
-
   @classmethod
   def setup_class(cls):
-
     if not is_live_cluster():
       raise SkipTest('Sentry tests require a live sentry server')
 
@@ -61,9 +59,26 @@ class TestWithSentry:
     cls.config_path = os.path.join(SENTRY_CONF_DIR.get(), 'sentry-site.xml')
 
 
-  @classmethod
-  def teardown_class(cls):
+  def setUp(self):
+    self.rpc_addresses = ''
+    if sentry_site.get_sentry_server_rpc_addresses() is not None:
+      self.rpc_addresses = ','.join(sentry_site.get_sentry_server_rpc_addresses())
+
+    self.tmpdir = tempfile.mkdtemp()
+    self.resets = [
+      SENTRY_CONF_DIR.set_for_testing(self.tmpdir),
+    ]
+    if API_CACHE is not None:
+      self.resets.append(API_CACHE.set_for_testing(None))
+    if API2_CACHE is not None:
+      self.resets.append(API2_CACHE.set_for_testing(None))
+
+
+  def tearDown(self):
     sentry_site.reset()
+    for reset in self.resets:
+      reset()
+    shutil.rmtree(self.tmpdir)
 
 
   def test_get_collections(self):
@@ -75,84 +90,83 @@ class TestWithSentry:
     assert_equal(0, resp.status.value, resp)
 
 
-  def test_ha_failover(self):
-    try:
-      tmpdir = tempfile.mkdtemp()
-      finish = SENTRY_CONF_DIR.set_for_testing(tmpdir)
-      rpc_addresses = ','.join(sentry_site.get_sentry_server_rpc_addresses())
-
-      # Test with good-host,bad-host-1,bad-host-2
-      xml = self._sentry_site_xml(rpc_addresses='%s,bad-host-1,bad-host-2' % rpc_addresses)
-      file(os.path.join(tmpdir, 'sentry-site.xml'), 'w').write(xml)
-      sentry_site.reset()
-
-      api = get_api(self.user)
-      assert_equal('%s,bad-host-1,bad-host-2' % rpc_addresses, ','.join(sentry_site.get_sentry_server_rpc_addresses()))
-      resp = api.list_sentry_roles_by_group(groupName='*')
-      assert_true(isinstance(resp, list))
-
-      api2 = get_api2(self.user, 'solr')
-      resp = api2.list_sentry_roles_by_group(groupName='*')
-      assert_true(isinstance(resp, list))
-
-      # Test with bad-host-1,bad-host-2,good-host
-      xml = self._sentry_site_xml(rpc_addresses='bad-host-1,bad-host-2,%s' % rpc_addresses)
-      file(os.path.join(tmpdir, 'sentry-site.xml'), 'w').write(xml)
-      sentry_site.reset()
-
-      api = get_api(self.user)
-      assert_equal('bad-host-1,bad-host-2,%s' % rpc_addresses, ','.join(sentry_site.get_sentry_server_rpc_addresses()))
-      resp = api.list_sentry_roles_by_group(groupName='*')
-      assert_true(isinstance(resp, list))
-
-      api2 = get_api2(self.user, 'solr')
-      resp = api2.list_sentry_roles_by_group(groupName='*')
-      assert_true(isinstance(resp, list))
-
-      # Test with bad-host-1,good-host,bad-host-2
-      xml = self._sentry_site_xml(rpc_addresses='bad-host-1,%s,bad-host-2' % rpc_addresses)
-      file(os.path.join(tmpdir, 'sentry-site.xml'), 'w').write(xml)
-      sentry_site.reset()
-
-      api = get_api(self.user)
-      assert_equal('bad-host-1,%s,bad-host-2' % rpc_addresses, ','.join(sentry_site.get_sentry_server_rpc_addresses()))
-      resp = api.list_sentry_roles_by_group(groupName='*')
-      assert_true(isinstance(resp, list))
-
-      api2 = get_api2(self.user, 'solr')
-      resp = api2.list_sentry_roles_by_group(groupName='*')
-      assert_true(isinstance(resp, list))
-
-      # Test with all bad hosts
-      xml = self._sentry_site_xml(rpc_addresses='bad-host-1,bad-host-2')
-      file(os.path.join(tmpdir, 'sentry-site.xml'), 'w').write(xml)
-      sentry_site.reset()
-
-      api = get_api(self.user)
-      assert_equal('bad-host-1,bad-host-2', ','.join(sentry_site.get_sentry_server_rpc_addresses()))
-      assert_raises(PopupException, api.list_sentry_roles_by_group, '*')
-
-      api2 = get_api2(self.user, 'solr')
-      assert_raises(PopupException, api2.list_sentry_roles_by_group, '*')
-
-      # Test with no rpc hosts and fallback to hostname and port
-      xml = self._sentry_site_xml(rpc_addresses='')
-      file(os.path.join(tmpdir, 'sentry-site.xml'), 'w').write(xml)
-      sentry_site.reset()
-
-      api = get_api(self.user)
-      assert_false(sentry_site.is_ha_enabled(), sentry_site.get_sentry_server_rpc_addresses())
-      assert_true(is_enabled() and HOSTNAME.get() and HOSTNAME.get() != 'localhost')
-      resp = api.list_sentry_roles_by_group(groupName='*')
-      assert_true(isinstance(resp, list))
-
-      api2 = get_api2(self.user, 'solr')
-      resp = api2.list_sentry_roles_by_group(groupName='*')
-      assert_true(isinstance(resp, list))
-    finally:
-      sentry_site.reset()
-      finish()
-      shutil.rmtree(tmpdir)
+  def test_ha_failover_good_bad_bad(self):
+    # Test with good-host,bad-host-1,bad-host-2
+    xml = self._sentry_site_xml(rpc_addresses='%s,bad-host-1,bad-host-2' % self.rpc_addresses)
+    file(os.path.join(self.tmpdir, 'sentry-site.xml'), 'w').write(xml)
+    sentry_site.reset()
+
+    api = get_api(self.user)
+    assert_equal('%s,bad-host-1,bad-host-2' % self.rpc_addresses, ','.join(sentry_site.get_sentry_server_rpc_addresses()))
+    resp = api.list_sentry_roles_by_group(groupName='*')
+    assert_true(isinstance(resp, list))
+
+    api2 = get_api2(self.user, 'solr')
+    resp = api2.list_sentry_roles_by_group(groupName='*')
+    assert_true(isinstance(resp, list))
+
+
+  def test_ha_failover_bad_bad_good(self):
+    # Test with bad-host-1,bad-host-2,good-host
+    xml = self._sentry_site_xml(rpc_addresses='bad-host-1,bad-host-2,%s' % self.rpc_addresses)
+    file(os.path.join(self.tmpdir, 'sentry-site.xml'), 'w').write(xml)
+    sentry_site.reset()
+
+    api = get_api(self.user)
+    assert_equal('bad-host-1,bad-host-2,%s' % self.rpc_addresses, ','.join(sentry_site.get_sentry_server_rpc_addresses()))
+    resp = api.list_sentry_roles_by_group(groupName='*')
+    assert_true(isinstance(resp, list))
+
+    api2 = get_api2(self.user, 'solr')
+    resp = api2.list_sentry_roles_by_group(groupName='*')
+    assert_true(isinstance(resp, list))
+
+
+  def test_ha_failover_bad_good_bad(self):
+    # Test with bad-host-1,good-host,bad-host-2
+    xml = self._sentry_site_xml(rpc_addresses='bad-host-1,%s,bad-host-2' % self.rpc_addresses)
+    file(os.path.join(self.tmpdir, 'sentry-site.xml'), 'w').write(xml)
+    sentry_site.reset()
+
+    api = get_api(self.user)
+    assert_equal('bad-host-1,%s,bad-host-2' % self.rpc_addresses, ','.join(sentry_site.get_sentry_server_rpc_addresses()))
+    resp = api.list_sentry_roles_by_group(groupName='*')
+    assert_true(isinstance(resp, list))
+
+    api2 = get_api2(self.user, 'solr')
+    resp = api2.list_sentry_roles_by_group(groupName='*')
+    assert_true(isinstance(resp, list))
+
+
+  def test_ha_failover_all_bad(self):
+    # Test with all bad hosts
+    xml = self._sentry_site_xml(rpc_addresses='bad-host-1,bad-host-2')
+    file(os.path.join(self.tmpdir, 'sentry-site.xml'), 'w').write(xml)
+    sentry_site.reset()
+
+    api = get_api(self.user)
+    assert_equal('bad-host-1,bad-host-2', ','.join(sentry_site.get_sentry_server_rpc_addresses()))
+    assert_raises(PopupException, api.list_sentry_roles_by_group, '*')
+
+    api2 = get_api2(self.user, 'solr')
+    assert_raises(PopupException, api2.list_sentry_roles_by_group, '*')
+
+
+  def test_no_rpc_hosts(self):
+    # Test with no rpc hosts and fallback to hostname and port
+    xml = self._sentry_site_xml(rpc_addresses='')
+    file(os.path.join(self.tmpdir, 'sentry-site.xml'), 'w').write(xml)
+    sentry_site.reset()
+
+    api = get_api(self.user)
+    assert_false(sentry_site.is_ha_enabled(), sentry_site.get_sentry_server_rpc_addresses())
+    assert_true(is_enabled() and HOSTNAME.get() and HOSTNAME.get() != 'localhost')
+    resp = api.list_sentry_roles_by_group(groupName='*')
+    assert_true(isinstance(resp, list))
+
+    api2 = get_api2(self.user, 'solr')
+    resp = api2.list_sentry_roles_by_group(groupName='*')
+    assert_true(isinstance(resp, list))
 
 
   def _sentry_site_xml(self, rpc_addresses):