jina-ai · JoanFM · May 12, 2022 · May 6, 2022 · May 9, 2022 · May 10, 2022
@@ -83,6 +83,7 @@ def _warning_on_one_line(message, category, filename, lineno, *args, **kwargs):
 # 2. grep -rohEI --exclude-dir=jina/hub --exclude-dir=tests --include \*.py "\'JINA_.*?\'" jina  | sort -u | sed "s/$/,/g"
 # 3. copy all lines EXCEPT the first (which is the grep command in the last line)
 __jina_env__ = (
+    'JINA_AUTH_TOKEN',
     'JINA_DEFAULT_HOST',
     'JINA_DEFAULT_TIMEOUT_CTRL',
     'JINA_DEFAULT_WORKSPACE_BASE',

diff --git a/jina/hubble/helper.py b/jina/hubble/helper.py
@@ -20,6 +20,7 @@
 
 from jina import __resources_path__
 from jina.enums import BetterEnum
+from jina.helper import get_request_header as _get_request_header_main
 from jina.importer import ImportExtensions
 from jina.logging.predefined import default_logger
 
@@ -34,6 +35,16 @@ def _get_hub_root() -> Path:
     return hub_root
 
 
+@lru_cache()
+def _get_hub_config() -> Optional[Dict]:
+    hub_root = _get_hub_root()
+
+    config_file = hub_root.joinpath('config.json')
+    if config_file.exists():
+        with open(config_file) as f:
+            return json.load(f)
+
+
 @lru_cache()
 def get_hub_packages_dir() -> Path:
     """Get the path of folder where the hub packages are stored
@@ -106,6 +117,37 @@ def _get_hubble_base_url() -> str:
     return u
 
 
+@lru_cache()
+def _get_auth_token() -> Optional[str]:
+    """Get user auth token.
+    .. note:: We first check `JINA_AUTH_TOKEN` environment variable.
+        if token is not None, use env token. Otherwise, we get token from config.
+
+    :return: user auth token
+    """
+    token_from_env = os.environ.get('JINA_AUTH_TOKEN')
+    if token_from_env:
+        return token_from_env
+
+    config = _get_hub_config()
+    if config:
+        return config.get('auth_token')
+
+
+def get_request_header() -> Dict:
+    """Return the header of request with an authorization token.
+
+    :return: request header
+    """
+    headers = _get_request_header_main()
+
+    auth_token = _get_auth_token()
+    if auth_token:
+        headers['Authorization'] = f'token {auth_token}'
+
+    return headers
+
+
 def get_hubble_url_v1() -> str:
     """Get v1 Hubble Url
 

diff --git a/jina/hubble/hubio.py b/jina/hubble/hubio.py
@@ -10,13 +10,7 @@
 from typing import Dict, Optional, Union
 
 from jina import __resources_path__, __version__
-from jina.helper import (
-    ArgNamespace,
-    colored,
-    get_request_header,
-    get_rich_console,
-    retry,
-)
+from jina.helper import ArgNamespace, get_rich_console, retry
 from jina.hubble import HubExecutor
 from jina.hubble.helper import (
     archive_package,
@@ -26,6 +20,7 @@
     get_download_cache_dir,
     get_hubble_error_message,
     get_hubble_url_v2,
+    get_request_header,
     parse_hub_uri,
     upload_file,
 )
@@ -393,7 +388,7 @@ def push(self) -> None:
                     form_data['secret'] = self.args.secret or secret
 
                 st.update(f'Connecting to Jina Hub ...')
-                if form_data.get('id') and form_data.get('secret'):
+                if form_data.get('id'):
                     hubble_url = get_hubble_url_v2() + '/rpc/executor.update'
                 else:
                     hubble_url = get_hubble_url_v2() + '/rpc/executor.create'
@@ -461,7 +456,7 @@ def push(self) -> None:
                 if image:
                     new_uuid8, new_secret = self._prettyprint_result(console, image)
                     if new_uuid8 != uuid8 or new_secret != secret:
-                        dump_secret(work_path, new_uuid8, new_secret)
+                        dump_secret(work_path, new_uuid8, new_secret or '')
                 else:
                     raise Exception(f'Unknown Error, session_id: {session_id}')
 
@@ -481,7 +476,7 @@ def _prettyprint_result(self, console, image):
         from rich.table import Table
 
         uuid8 = image['id']
-        secret = image['secret']
+        secret = image.get('secret')
         visibility = image['visibility']
         tag = self.args.tag[0] if self.args.tag else None
 
@@ -494,11 +489,14 @@ def _prettyprint_result(self, console, image):
         )
         if 'name' in image:
             table.add_row(':name_badge: Name', image['name'])
-        table.add_row(':lock: Secret', secret)
-        table.add_row(
-            '',
-            ':point_up:️ [bold red]Please keep this token in a safe place!',
-        )
+
+        if secret:
+            table.add_row(':lock: Secret', secret)
+            table.add_row(
+                '',
+                ':point_up:️ [bold red]Please keep this token in a safe place!',
+            )
+
         table.add_row(':eyes: Visibility', visibility)
 
         p1 = Panel(
@@ -511,7 +509,9 @@ def _prettyprint_result(self, console, image):
 
         presented_id = image.get('name', uuid8)
         usage = (
-            f'{presented_id}' if visibility == 'public' else f'{presented_id}:{secret}'
+            f'{presented_id}'
+            if visibility == 'public' or not secret
+            else f'{presented_id}:{secret}'
         ) + (f'/{tag}' if tag else '')
 
         if not self.args.no_usage:
@@ -789,7 +789,7 @@ def pull(self) -> str:
                 presented_id = getattr(executor, 'name', executor.uuid)
                 executor_name = (
                     f'{presented_id}'
-                    if executor.visibility == 'public'
+                    if executor.visibility == 'public' or not secret
                     else f'{presented_id}:{secret}'
                 ) + (f'/{tag}' if tag else '')
 

diff --git a/tests/unit/hubble/test_hubio.py b/tests/unit/hubble/test_hubio.py
@@ -13,7 +13,12 @@
 import requests
 import yaml
 
-from jina.hubble.helper import disk_cache_offline
+from jina.hubble.helper import (
+    _get_auth_token,
+    _get_hub_config,
+    _get_hub_root,
+    disk_cache_offline,
+)
 from jina.hubble.hubio import HubExecutor, HubIO
 from jina.parsers.hubble import (
     set_hub_new_parser,
@@ -24,6 +29,27 @@
 cur_dir = os.path.dirname(os.path.abspath(__file__))
 
 
+def clear_function_caches():
+    _get_auth_token.cache_clear()
+    _get_hub_root.cache_clear()
+    _get_hub_config.cache_clear()
+
+
+@pytest.fixture(scope='function')
+def auth_token(tmpdir):
+    clear_function_caches()
+    os.environ['JINA_HUB_ROOT'] = str(tmpdir)
+
+    token = 'test-auth-token'
+    with open(tmpdir / 'config.json', 'w') as f:
+        json.dump({'auth_token': token}, f)
+
+    yield token
+
+    clear_function_caches()
+    del os.environ['JINA_HUB_ROOT']
+
+
 class PostMockResponse:
     def __init__(self, response_code: int = 201):
         self.response_code = response_code
@@ -212,6 +238,31 @@ def _mock_post(url, data, headers=None, stream=True):
     )
 
 
+def test_push_with_authorization(mocker, monkeypatch, auth_token):
+    mock = mocker.Mock()
+
+    def _mock_post(url, data, headers, stream):
+        mock(url=url, headers=headers)
+        return PostMockResponse(response_code=200)
+
+    monkeypatch.setattr(requests, 'post', _mock_post)
+
+    exec_path = os.path.join(cur_dir, 'dummy_executor')
+    args = set_hub_push_parser().parse_args([exec_path])
+
+    HubIO(args).push()
+
+    # remove .jina
+    exec_config_path = os.path.join(exec_path, '.jina')
+    shutil.rmtree(exec_config_path)
+
    assert mock.call_count == 1
+
+    _, kwargs = mock.call_args_list[0]
+
+    assert kwargs['headers'].get('Authorization') == f'token {auth_token}'
+
+
 @pytest.mark.parametrize('rebuild_image', [True, False])
 def test_fetch(mocker, monkeypatch, rebuild_image):
     mock = mocker.Mock()
@@ -293,6 +344,24 @@ def _mock_post(url, json, headers=None):
     assert mock.call_count == 6  # mock must be called 3+3
 
 
+def test_fetch_with_authorization(mocker, monkeypatch, auth_token):
+    mock = mocker.Mock()
+
+    def _mock_post(url, json, headers):
+        mock(url=url, json=json, headers=headers)
+        return FetchMetaMockResponse(response_code=200)
+
+    monkeypatch.setattr(requests, 'post', _mock_post)
+
+    HubIO.fetch_meta('dummy_mwu_encoder', tag=None, force=True)
+
+    assert mock.call_count == 1
+
+    _, kwargs = mock.call_args_list[0]
+
+    assert kwargs['headers'].get('Authorization') == f'token {auth_token}'
+
+
 class DownloadMockResponse:
     def __init__(self, response_code: int = 200):
         self.response_code = response_code