feat(login): api key (#840)

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

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,15 +63,15 @@ 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
         # if url is not Routes.login and self.token:
         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
 

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:

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()

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. <[email protected]>"]
 readme = "README.md"

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)

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())

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"

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"

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",