Add support for api key (#119)

* Add support for api key

Signed-off-by: wslulciuc <[email protected]>

marquez_client/backend.py

@@ -10,6 +10,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 class Backend:
     def put(self, path, headers, json):
         pass

marquez_client/client.py

@@ -17,28 +17,27 @@
 from six.moves.urllib.parse import quote
 
 from marquez_client import errors
-from marquez_client.constants import (DEFAULT_TIMEOUT_MS)
+from marquez_client.constants import (DEFAULT_TIMEOUT_MS, API_PATH_V1)
 from marquez_client.models import (DatasetType, JobType)
 from marquez_client.utils import Utils
 from marquez_client.version import VERSION
 
-_API_PATH = '/api/v1'
 _USER_AGENT = f'marquez-python/{VERSION}'
 _HEADERS = {'User-Agent': _USER_AGENT}
 
 log = logging.getLogger(__name__)
 
 
 # Marquez Client
-class MarquezClient(object):
-    def __init__(self, url, timeout_ms=None):
+class MarquezClient:
+    def __init__(self, url, timeout_ms=None, api_key: str = None):
         self._timeout = Utils.to_seconds(timeout_ms or os.environ.get(
             'MARQUEZ_TIMEOUT_MS', DEFAULT_TIMEOUT_MS)
         )
+        self._api_base = f"{url}{API_PATH_V1}"
 
-        self._api_base = f'{url}{_API_PATH}'
-
-        log.debug(self._api_base)
+        if api_key:
+            Utils.add_auth_to(_HEADERS, api_key)
 
     # Namespace API
     def create_namespace(self, namespace_name, owner_name, description=None):

marquez_client/clients.py

@@ -10,48 +10,49 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
 import os
 
-import marquez_client
-from marquez_client.constants import (DEFAULT_MARQUEZ_BACKEND,
-                                      DEFAULT_MARQUEZ_URL,
-                                      DEFAULT_MARQUEZ_FILE,
-                                      DEFAULT_TIMEOUT_MS)
-from marquez_client.file_backend import FileBackend
+from marquez_client import MarquezClient, MarquezWriteOnlyClient
 from marquez_client.http_backend import HttpBackend
+from marquez_client.file_backend import FileBackend
 from marquez_client.log_backend import LogBackend
 from marquez_client.utils import Utils
-
-log = logging.getLogger(__name__)
+from marquez_client.constants import (
+    DEFAULT_MARQUEZ_BACKEND,
+    DEFAULT_MARQUEZ_URL,
+    DEFAULT_MARQUEZ_FILE,
+    DEFAULT_TIMEOUT_MS
+)
 
 
 # Marquez Clients
 class Clients(object):
-    def __init__(self):
-        log.debug("Clients.init")
-
     @staticmethod
     def new_client():
-        url = os.environ.get('MARQUEZ_URL', DEFAULT_MARQUEZ_URL)
-        return marquez_client.MarquezClient(url)
+        return MarquezClient(
+            url=os.environ.get('MARQUEZ_URL', DEFAULT_MARQUEZ_URL),
+            api_key=os.environ.get('MARQUEZ_API_KEY')
+        )
 
     @staticmethod
     def new_write_only_client():
-        return marquez_client.MarquezWriteOnlyClient(Clients.from_env())
+        return MarquezWriteOnlyClient(
+            backend=Clients._backend_from_env(),
+        )
 
     @staticmethod
-    def from_env():
-        backend_env = \
+    def _backend_from_env():
+        backend = \
             os.environ.get('MARQUEZ_BACKEND', DEFAULT_MARQUEZ_BACKEND).upper()
 
-        if backend_env == 'HTTP':
+        if backend == 'HTTP':
             url = os.environ.get('MARQUEZ_URL', DEFAULT_MARQUEZ_URL)
+            api_key = os.environ.get('MARQUEZ_API_KEY')
             timeout = Utils.to_seconds(
                 os.environ.get('MARQUEZ_TIMEOUT_MS', DEFAULT_TIMEOUT_MS))
-            return HttpBackend(url, timeout)
-        elif backend_env == 'FILE':
+            return HttpBackend(url, timeout, api_key)
+        elif backend == 'FILE':
             file = os.environ.get('MARQUEZ_FILE', DEFAULT_MARQUEZ_FILE)
             return FileBackend(file)
-        elif backend_env == 'LOG':
+        elif backend == 'LOG':
             return LogBackend()

marquez_client/constants.py

@@ -15,3 +15,5 @@
 DEFAULT_MARQUEZ_BACKEND = 'http'
 DEFAULT_MARQUEZ_URL = 'http://localhost:8080'
 DEFAULT_MARQUEZ_FILE = '/tmp/marquez/client.requests.log'
+
+API_PATH_V1 = '/api/v1'

marquez_client/http_backend.py

@@ -10,36 +10,43 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
-
 import requests
 
 from marquez_client import errors
 from marquez_client.backend import Backend
-
-log = logging.getLogger(__name__)
+from marquez_client.utils import Utils
+from marquez_client.constants import API_PATH_V1
 
 
 class HttpBackend(Backend):
-    def __init__(self, url, timeout):
+    def __init__(self, url, timeout, api_key: str = None):
         self._timeout = timeout
-        self._url = url
+        self._api_base = f"{url}{API_PATH_V1}"
+        self._api_key = api_key
 
     def put(self, path, headers, payload):
-        log.debug("_put()")
+        if self._api_key:
+            Utils.add_auth_to(headers, self._api_key)
 
         response = requests.put(
-            url=f'{self._url}{path}', headers=headers, json=payload,
-            timeout=self._timeout)
+            url=f"{self._api_base}{path}",
+            headers=headers,
+            json=payload,
+            timeout=self._timeout
+        )
 
         return self._response(response, as_json=True)
 
     def post(self, path, headers, payload=None):
-        log.debug("_post()")
+        if self._api_key:
+            Utils.add_auth_to(headers, self._api_key)
 
         response = requests.post(
-            url=f'{self._url}{path}', headers=headers, json=payload,
-            timeout=self._timeout)
+            url=f"{self._api_base}{path}",
+            headers=headers,
+            json=payload,
+            timeout=self._timeout
+        )
 
         return self._response(response, as_json=True)
 

marquez_client/utils.py

@@ -99,3 +99,7 @@ def utc_now():
     def get_json(file):
         with open(file) as json_file:
             return json.load(json_file)
+
+    @staticmethod
+    def add_auth_to(headers, api_key):
+        headers['Authorization'] = f"Bearer: {api_key}"

tests/test_marquez_clients.py

@@ -0,0 +1,70 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+from marquez_client import Clients
+from marquez_client.http_backend import HttpBackend
+from marquez_client.file_backend import FileBackend
+from marquez_client.log_backend import LogBackend
+from marquez_client.utils import Utils
+from marquez_client.constants import DEFAULT_MARQUEZ_URL, API_PATH_V1
+
+API_KEY = 'PuRx8GT3huSXlheDIRUK1YUatGpLVEuL'
+API_BASE = f"{DEFAULT_MARQUEZ_URL}{API_PATH_V1}"
+
+
+def test_new_client():
+    os.environ['MARQUEZ_API_KEY'] = API_KEY
+
+    from marquez_client.client import _USER_AGENT, _HEADERS
+    headers_with_auth = {'User-Agent': _USER_AGENT}
+
+    # Add API key to headers
+    Utils.add_auth_to(headers_with_auth, API_KEY)
+
+    client = Clients.new_client()
+    assert client._api_base == API_BASE
+    assert _HEADERS == headers_with_auth
+
+    del os.environ['MARQUEZ_API_KEY']
+
+
+def test_new_wo_client_http():
+    os.environ['MARQUEZ_BACKEND'] = 'http'
+    os.environ['MARQUEZ_API_KEY'] = API_KEY
+
+    client = Clients.new_write_only_client()
+    assert isinstance(client._backend, HttpBackend)
+    assert client._backend._api_base == API_BASE
+    assert client._backend._api_key == API_KEY
+
+    del os.environ['MARQUEZ_BACKEND']
+    del os.environ['MARQUEZ_API_KEY']
+
+
+def test_new_wo_client_file():
+    os.environ['MARQUEZ_BACKEND'] = 'file'
+
+    client = Clients.new_write_only_client()
+    assert isinstance(client._backend, FileBackend)
+
+    del os.environ['MARQUEZ_BACKEND']
+
+
+def test_new_wo_client_log():
+    os.environ['MARQUEZ_BACKEND'] = 'log'
+
+    client = Clients.new_write_only_client()
+    assert isinstance(client._backend, LogBackend)
+
+    del os.environ['MARQUEZ_BACKEND']