[FEATURE] Add api key login (#384)

* Add api key login

* revert notebook

* try to open the browser for them if possible

* Direct to password reset page, minor cleanup

* Fix bad credentials test

* token hidden

* Update __init__.py

Co-authored-by: Ben Epstein <[email protected]>

dataquality/__init__.py

@@ -1,6 +1,6 @@
 "dataquality"
 
-__version__ = "v0.5.5"
+__version__ = "v0.6.0"
 
 import os
 import resource
@@ -9,7 +9,7 @@
 import dataquality.integrations
 import dataquality.metrics
 from dataquality.core._config import config
-from dataquality.core.auth import login
+from dataquality.core.auth import login, logout
 from dataquality.core.finish import finish, get_run_status, wait_for_run
 from dataquality.core.init import init
 from dataquality.core.log import (
@@ -65,6 +65,7 @@ def configure() -> None:
 __all__ = [
     "__version__",
     "login",
+    "logout",
     "init",
     "log_data_samples",
     "log_model_outputs",

dataquality/core/_config.py

@@ -1,7 +1,7 @@
 import json
 import os
 import warnings
-from enum import Enum, unique
+from enum import Enum
 from pathlib import Path
 from typing import Dict, Optional
 
@@ -45,14 +45,8 @@ class ConfigData(str, Enum):
     DEFAULT_GALILEO_CONFIG_FILE = f"{DEFAULT_GALILEO_CONFIG_DIR}/config.json"
 
 
-@unique
-class AuthMethod(str, Enum):
-    email = "email"
-
-
 class Config(BaseModel):
     api_url: str
-    auth_method: AuthMethod = AuthMethod.email
     token: Optional[str] = None
     current_user: Optional[str] = None
     current_project_id: Optional[UUID4] = None

dataquality/core/auth.py

@@ -1,11 +1,11 @@
 import getpass
 import os
-from typing import Callable, Dict
+import webbrowser
 
 import requests
 
 from dataquality.clients.api import ApiClient
-from dataquality.core._config import AuthMethod, Config, config
+from dataquality.core._config import config
 from dataquality.exceptions import GalileoException
 from dataquality.utils.helpers import check_noop
 
@@ -14,36 +14,41 @@
 
 
 class _Auth:
-    def auth_methods(self) -> Dict[AuthMethod, Callable]:
-        return {AuthMethod.email: self.email_login}
-
     def email_login(self) -> None:
-        if self.email_token_present_and_valid(config):
-            return
-
         username = os.getenv("GALILEO_USERNAME")
         password = os.getenv("GALILEO_PASSWORD")
-        if not username or not password:
-            username = input("📧 Enter your email:")
-            password = getpass.getpass("🤫 Enter your password:")
         res = requests.post(
             f"{config.api_url}/login",
             data={
                 "username": username,
                 "password": password,
-                "auth_method": config.auth_method,
+                "auth_method": "email",
             },
             headers={"X-Galileo-Request-Source": "dataquality_python_client"},
         )
         if res.status_code != 200:
-            raise GalileoException(f"Issue authenticating: {res.json()['detail']}")
+            raise GalileoException(
+                (
+                    f"Issue authenticating: {res.json()['detail']} "
+                    "If you need to reset your password, "
+                    f"go to {config.api_url.replace('api', 'console')}/forgot-password"
+                )
+            )
 
         access_token = res.json().get("access_token")
         config.token = access_token
         config.update_file_config()
 
-    def email_token_present_and_valid(self, config: Config) -> bool:
-        return config.auth_method == "email" and api_client.valid_current_user()
+    def token_login(self) -> None:
+        token_url = config.api_url.replace("api.", "console.") + "/get-token"
+        try:
+            webbrowser.open(token_url)
+        except Exception:
+            pass
+        print(f"Go to {token_url} to generate a new API Key")
+        access_token = getpass.getpass("🔐 Enter your API Key:")
+        config.token = access_token
+        config.update_file_config()
 
 
 @check_noop
@@ -56,40 +61,32 @@ def login() -> None:
     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
     """
-    print(f"📡 {config.api_url.replace('api.','console.')}")
+    if api_client.valid_current_user():
+        print(f"✅ Already logged in as {config.current_user}!")
+        print("Use logout() if you want to change users")
+
+        return
+
+    print(f"📡 {config.api_url.replace('api','console')}")
     print("🔭 Logging you into Galileo\n")
-    auth_methods = ",".join([am.value for am in AuthMethod])
-    # Try auto auth config
-    if len(list(AuthMethod)) == 1:
-        auth_method = list(AuthMethod)[0].value
-        os.environ[GALILEO_AUTH_METHOD] = auth_method
-    auth_method = os.getenv(GALILEO_AUTH_METHOD) or config.auth_method
-    if not auth_method or auth_method.lower() not in list(AuthMethod):
-        # We currently only have 1 auth method, so no reason to ask the user
-        if len(list(AuthMethod)) == 1:
-            auth_method = list(AuthMethod)[0].value
-        else:
-            auth_method = input(
-                "🔐 How would you like to login? \n"
-                f"Enter one of the following: {auth_methods}\n"
-            )
-        if auth_method.lower() not in list(AuthMethod):
-            print(
-                "Invalid login request. You must input one of "
-                f"the following authentication methods: {auth_methods}."
-            )
-            return
-        else:  # Save it as an environment variable for the next login
-            print("🤝 Saving preferred login method")
-            os.environ[GALILEO_AUTH_METHOD] = auth_method
-    else:
-        print(f"👀 Found auth method {auth_method} set via env, skipping prompt.")
-    config.auth_method = AuthMethod(auth_method)
+
     _auth = _Auth()
-    _auth.auth_methods()[config.auth_method]()
+    if os.getenv("GALILEO_USERNAME") and os.getenv("GALILEO_PASSWORD"):
+        _auth.email_login()
+    else:
+        _auth.token_login()
+
     current_user_email = api_client.get_current_user().get("email")
     if not current_user_email:
         return
     config.current_user = current_user_email
     config.update_file_config()
     print(f"🚀 You're logged in to Galileo as {current_user_email}!")
+
+
+@check_noop
+def logout() -> None:
+    config.current_user = None
+    config.token = None
+    config.update_file_config()
+    print("👋 You have logged out of Galileo")

tests/test_auth.py

@@ -33,4 +33,7 @@ def test_bad_login(mock_post: MagicMock, set_test_config: Callable) -> None:
     os.environ["GALILEO_PASSWORD"] = "password"
     with pytest.raises(GalileoException) as e:
         dataquality.login()
-    assert e.value.args[0] == "Issue authenticating: Incorrect login credentials."
+    assert e.value.args[0] == (
+        "Issue authenticating: Incorrect login credentials. If you need to reset your "
+        "password, go to http://localhost:8088/forgot-password"
+    )