From 980b641eaa4a3ea788ed4cbe4b01d4b7da92186e Mon Sep 17 00:00:00 2001 From: hunhoon21 Date: Tue, 6 Aug 2024 17:29:05 +0900 Subject: [PATCH] :hammer: add HTTPBasicAuth for supporting MLFlow authentication --- mlflow_export_import/client/http_client.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/mlflow_export_import/client/http_client.py b/mlflow_export_import/client/http_client.py index 475bb29..a92900f 100644 --- a/mlflow_export_import/client/http_client.py +++ b/mlflow_export_import/client/http_client.py @@ -1,6 +1,7 @@ from abc import abstractmethod, ABCMeta import os import json +from requests.auth import HTTPBasicAuth import requests import click from mlflow_export_import.common import MlflowExportImportException @@ -96,11 +97,14 @@ def __init__(self, api_name, host=None, token=None): self.host = host self.api_uri = os.path.join(host, api_name) self.token = token + user = os.environ.get("MLFLOW_TRACKING_USERNAME") + password = os.environ.get("MLFLOW_TRACKING_PASSWORD") + self.auth = HTTPBasicAuth(user, password) if user and password else None def _get(self, resource, params=None): uri = self._mk_uri(resource) - rsp = requests.get(uri, headers=self._mk_headers(), data=params, timeout=_TIMEOUT) + rsp = requests.get(uri, headers=self._mk_headers(), data=params, auth=self.auth, timeout=_TIMEOUT) return self._check_response(rsp, params) @@ -151,7 +155,7 @@ def patch(self, resource, data=None): def _delete(self, resource): uri = self._mk_uri(resource) - rsp = requests.delete(uri, headers=self._mk_headers(), timeout=_TIMEOUT) + rsp = requests.delete(uri, headers=self._mk_headers(), auth=self.auth, timeout=_TIMEOUT) return self._check_response(rsp) def delete(self, resource): @@ -163,7 +167,7 @@ def delete(self, resource): def _mutator(self, method, resource, data=None): uri = self._mk_uri(resource) - rsp = method(uri, headers=self._mk_headers(), data=data, timeout=_TIMEOUT) + rsp = method(uri, headers=self._mk_headers(), data=data, auth=self.auth, timeout=_TIMEOUT) return self._check_response(rsp)