diff --git a/marquez_client/backend.py b/marquez_client/backend.py index 8658cb7..0fe224d 100644 --- a/marquez_client/backend.py +++ b/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 diff --git a/marquez_client/client.py b/marquez_client/client.py index 4ef4489..888008b 100644 --- a/marquez_client/client.py +++ b/marquez_client/client.py @@ -17,12 +17,11 @@ 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} @@ -30,15 +29,15 @@ # 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): diff --git a/marquez_client/clients.py b/marquez_client/clients.py index ad1799f..df7f237 100644 --- a/marquez_client/clients.py +++ b/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() diff --git a/marquez_client/constants.py b/marquez_client/constants.py index 94983be..4e9620f 100644 --- a/marquez_client/constants.py +++ b/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' diff --git a/marquez_client/http_backend.py b/marquez_client/http_backend.py index 6ba57f2..7c452c5 100644 --- a/marquez_client/http_backend.py +++ b/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) diff --git a/marquez_client/utils.py b/marquez_client/utils.py index e17c66d..15bf5bb 100644 --- a/marquez_client/utils.py +++ b/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}" diff --git a/tests/test_marquez_clients.py b/tests/test_marquez_clients.py new file mode 100644 index 0000000..3041273 --- /dev/null +++ b/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']