diff --git a/dataquality/__init__.py b/dataquality/__init__.py index c26a326e8..a5070c35b 100644 --- a/dataquality/__init__.py +++ b/dataquality/__init__.py @@ -30,7 +30,7 @@ dataquality.get_insights() """ -__version__ = "1.6.1" +__version__ = "2.0.1" import sys from typing import Any, List, Optional diff --git a/dataquality/clients/api.py b/dataquality/clients/api.py index d34aa8c64..edf576519 100644 --- a/dataquality/clients/api.py +++ b/dataquality/clients/api.py @@ -21,24 +21,31 @@ class ApiClient: - def _refresh_token(self) -> str: + def _refresh_jwt_token(self) -> str: username = os.getenv("GALILEO_USERNAME") password = os.getenv("GALILEO_PASSWORD") - if username is None or password is None: + api_key = os.getenv("GALILEO_API_KEY") + + if api_key is not None: + res = requests.post( + f"{config.api_url}/login/api_key", + json={"api_key": api_key}, + ) + elif username is not None and password is not None: + res = requests.post( + f"{config.api_url}/login", + data={ + "username": username, + "password": password, + "auth_method": "email", + }, + ) + else: raise GalileoException( "You are not logged in. Call dataquality.login()\n" "GALILEO_USERNAME and GALILEO_PASSWORD must be set" ) - res = requests.post( - f"{config.api_url}/login", - data={ - "username": username, - "password": password, - "auth_method": "email", - }, - headers={"X-Galileo-Request-Source": "dataquality_python_client"}, - ) if res.status_code != 200: raise GalileoException( ( @@ -56,7 +63,7 @@ def _refresh_token(self) -> str: def get_token(self) -> str: token = config.token if not token: - token = self._refresh_token() + token = self._refresh_jwt_token() # Check to see if our token is expired before making a request # and refresh token if it's expired @@ -64,7 +71,7 @@ def get_token(self) -> str: if token: claims = jwt.decode(token, options={"verify_signature": False}) if claims.get("exp", 0) < time(): - token = self._refresh_token() + token = self._refresh_jwt_token() return token diff --git a/dataquality/core/__init__.py b/dataquality/core/__init__.py index 05cc23521..5177c61ff 100644 --- a/dataquality/core/__init__.py +++ b/dataquality/core/__init__.py @@ -21,6 +21,7 @@ def configure(do_login: bool = True, _internal: bool = False) -> None: * GALILEO_CONSOLE_URL * GALILEO_USERNAME * GALILEO_PASSWORD + * GALILEO_API_KEY """ a.log_function("dq/configure") if not _internal: diff --git a/dataquality/core/auth.py b/dataquality/core/auth.py index 4294ca5d6..b20f3293f 100644 --- a/dataquality/core/auth.py +++ b/dataquality/core/auth.py @@ -40,7 +40,8 @@ def login() -> None: access from the console. To skip the prompt for automated workflows, you can set `GALILEO_USERNAME` - (your email) and GALILEO_PASSWORD if you signed up with an email and password + (your email) and GALILEO_PASSWORD if you signed up with an email and password. + You can set `GALILEO_API_KEY` to your API key if you have one. """ if not config.api_url: updated_config = reset_config() @@ -54,7 +55,11 @@ def login() -> None: print("🔭 Logging you into Galileo\n") _auth = _Auth() - if os.getenv("GALILEO_USERNAME") and os.getenv("GALILEO_PASSWORD"): + has_api_key = os.getenv("GALILEO_API_KEY") + has_username_password = os.getenv("GALILEO_USERNAME") and os.getenv( + "GALILEO_PASSWORD" + ) + if has_api_key or has_username_password: _auth.login_with_env_vars() if not valid_current_user: _auth.login_with_token() diff --git a/pyproject.toml b/pyproject.toml index cfed59f7a..cdc61a72a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api" [tool.poetry] name = "dataquality" -version = "2.0.0" +version = "2.0.1" description = "" authors = ["Galileo Technologies, Inc. "] readme = "README.md" diff --git a/tests/clients/test_api.py b/tests/clients/test_api.py index 6b41ee049..aa6579d05 100644 --- a/tests/clients/test_api.py +++ b/tests/clients/test_api.py @@ -26,16 +26,16 @@ api_client = ApiClient() -@mock.patch.object(ApiClient, "_refresh_token") -def test_refresh_token( - mock_refresh_token: MagicMock, set_test_config: Callable +@mock.patch.object(ApiClient, "_refresh_jwt_token") +def test_refresh_jwt_token( + mock_refresh_jwt_token: MagicMock, set_test_config: Callable ) -> None: """Base case: Tests creating a new project and run""" set_test_config(token="expired-token") - mock_refresh_token.return_value = "my-token" + mock_refresh_jwt_token.return_value = "my-token" token = api_client.get_token() assert token == "my-token" - mock_refresh_token.assert_called_once() + mock_refresh_jwt_token.assert_called_once() @mock.patch("requests.get", side_effect=mocked_get_project_run) diff --git a/tests/conftest.py b/tests/conftest.py index f06ccc173..bd31bdcdb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -30,6 +30,8 @@ torch.set_default_device("cpu") except AttributeError: warnings.warn("Torch default device not set to CPU", GalileoWarning) + + DEFAULT_API_URL = "http://localhost:8088" UUID_STR = "399057bc-b276-4027-a5cf-48893ac45388" TEST_STORE_DIR = "TEST_STORE" @@ -97,6 +99,13 @@ def __init__( self.TEST_PATH = TEST_PATH +@pytest.fixture(autouse=True) +def set_scikitlearn_data_folder() -> None: + if os.environ.get("PYTEST_XDIST_WORKER_COUNT"): + pid = os.getpid() + os.environ["SCIKIT_LEARN_DATA"] = f"~/scikit_learn_data_{pid}" + + @pytest.fixture(scope="session") def test_session_vars() -> TestSessionVariables: pid = str(os.getpid()) diff --git a/tests/core/test_auth.py b/tests/core/test_auth.py index 720de81fc..d6c461cf1 100644 --- a/tests/core/test_auth.py +++ b/tests/core/test_auth.py @@ -18,7 +18,7 @@ @mock.patch("requests.post", side_effect=mocked_login_requests) @mock.patch("requests.get", side_effect=mocked_login_requests) -def test_good_login( +def test_login_username_password( mock_get_current_user: MagicMock, mock_login: MagicMock, set_test_config: Callable ) -> None: os.environ[GALILEO_AUTH_METHOD] = "email" @@ -28,8 +28,21 @@ def test_good_login( dataquality.login() +@mock.patch("requests.post", side_effect=mocked_login_requests) +@mock.patch("requests.get", side_effect=mocked_login_requests) +def test_login_api_key( + mock_get_current_user: MagicMock, mock_login: MagicMock, set_test_config: Callable +) -> None: + os.environ["GALILEO_API_KEY"] = "long-secret-key" + set_test_config(token="mytoken") + dataquality.login() + + @mock.patch("requests.post", side_effect=mocked_failed_login_requests) -def test_bad_login(mock_post: MagicMock, set_test_config: Callable) -> None: +def test_bad_login( + mock_post: MagicMock, set_test_config: Callable, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("GALILEO_API_KEY", raising=False) set_test_config(token=None) os.environ[GALILEO_AUTH_METHOD] = "email" os.environ["GALILEO_USERNAME"] = "user" diff --git a/tests/core/test_init.py b/tests/core/test_init.py index 4773a83e1..a1a2a8fff 100644 --- a/tests/core/test_init.py +++ b/tests/core/test_init.py @@ -603,13 +603,16 @@ def test_reconfigure_sets_env_vars(mock_login: MagicMock) -> None: assert mock_login.call_count == 2 +@patch("dataquality.core._config.input") @patch("requests.post", side_effect=mocked_login_requests) @patch("requests.get", side_effect=mocked_login_requests) def test_reconfigure_resets_user_token( mock_get_request: MagicMock, mock_post_request: MagicMock, + mock_input: MagicMock, set_test_config: Callable, ) -> None: + mock_input.return_value = "mock_url" set_test_config(token="old_token") os.environ[GALILEO_AUTH_METHOD] = "email" diff --git a/tests/test_utils/mock_request.py b/tests/test_utils/mock_request.py index 2f6ffaf45..854cfdb7d 100644 --- a/tests/test_utils/mock_request.py +++ b/tests/test_utils/mock_request.py @@ -51,7 +51,7 @@ def mocked_login_requests( timeout: Union[int, None] = None, files: Union[Dict, None] = None, ) -> MockResponse: - if request_url.endswith("login"): + if request_url.endswith("login") or request_url.endswith("api_key"): return MockResponse( { "access_token": "mock_token",