wip

Signed-off-by: byhsu <[email protected]>

flytekit/clients/auth/auth_client.py

@@ -3,7 +3,6 @@
 import base64 as _base64
 import hashlib as _hashlib
 import http.server as _BaseHTTPServer
-import logging
 import multiprocessing
 import os as _os
 import re as _re
@@ -20,6 +19,7 @@
 from .default_html import get_default_success_html
 from .exceptions import AccessTokenNotFoundError
 from .keyring import Credentials
+from flytekit import logger
 
 _code_verifier_length = 64
 _random_seed_length = 40
@@ -251,7 +251,7 @@ def _request_authorization_code(self):
         scheme, netloc, path, _, _, _ = _urlparse.urlparse(self._auth_endpoint)
         query = _urlencode(self._params)
         endpoint = _urlparse.urlunparse((scheme, netloc, path, None, query, None))
-        logging.debug(f"Requesting authorization code through {endpoint}")
+        logger.debug(f"Requesting authorization code through {endpoint}")
         _webbrowser.open_new_tab(endpoint)
 
     def _credentials_from_response(self, auth_token_resp) -> Credentials:

flytekit/clients/auth/authenticator.py

@@ -1,4 +1,3 @@
-import logging
 import subprocess
 import typing
 from abc import abstractmethod
@@ -9,6 +8,9 @@
 from .exceptions import AccessTokenNotFoundError, AuthenticationError
 from .keyring import Credentials, KeyringStore
 import click
+from flytekit import logger
+from rich.text import Text
+from rich.style import Style
 
 @dataclass
 class ClientConfig:
@@ -117,7 +119,7 @@ def refresh_credentials(self):
                     KeyringStore.store(self._creds)
                 return
             except AccessTokenNotFoundError:
-                logging.warning("Failed to refresh token. Kicking off a full authorization flow.")
+                logger.warning("Failed to refresh token. Kicking off a full authorization flow.")
                 KeyringStore.delete(self._endpoint)
 
         self._creds = self._auth_client.get_creds_from_remote()
@@ -140,11 +142,11 @@ def refresh_credentials(self):
         This function is used when the configuration value for AUTH_MODE is set to 'external_process'.
         It reads an id token generated by an external process started by running the 'command'.
         """
-        logging.debug("Starting external process to generate id token. Command {}".format(self._cmd))
+        logger.debug("Starting external process to generate id token. Command {}".format(self._cmd))
         try:
             output = subprocess.run(self._cmd, capture_output=True, text=True, check=True)
         except subprocess.CalledProcessError as e:
-            logging.error("Failed to generate token from command {}".format(self._cmd))
+            logger.error("Failed to generate token from command {}".format(self._cmd))
             raise AuthenticationError("Problems refreshing token with command: " + str(e))
         self._creds = Credentials(output.stdout.strip())
 
@@ -185,10 +187,10 @@ def refresh_credentials(self):
         scopes = self._scopes
 
         # Note that unlike the Pkce flow, the client ID does not come from Admin.
-        logging.debug(f"Basic authorization flow with client id {self._client_id} scope {scopes}")
+        logger.debug(f"Basic authorization flow with client id {self._client_id} scope {scopes}")
         authorization_header = token_client.get_basic_authorization_header(self._client_id, self._client_secret)
         token, expires_in = token_client.get_token(token_endpoint, scopes, authorization_header)
-        logging.info("Retrieved new token, expires in {}".format(expires_in))
+        logger.info("Retrieved new token, expires in {}".format(expires_in))
         self._creds = Credentials(token)
 
 
@@ -224,8 +226,14 @@ def __init__(
 
     def refresh_credentials(self):
         resp = token_client.get_device_code(self._device_auth_endpoint, self._client_id, self._audience, self._scope)
-        text = f"To Authenticate, navigate in a browser to the following URL: {click.style(resp.verification_uri, fg='blue', underline=True)} and enter code: {click.style(resp.user_code, fg='blue')}"
-        click.secho(text)
+        text = Text.assemble(
+            f"To Authenticate, navigate in a browser to the following URL: ",
+            Text(resp.verification_uri, style=Style(color="blue", underline=True)),
+            f" and enter code: ",
+            Text(resp.user_code, style=Style(color="blue")),
+        )
+        print("hello")
+        logger.info(text)
         try:
             # Currently the refresh token is not retreived. We may want to add support for refreshTokens so that
             # access tokens can be refreshed for once authenticated machines

flytekit/clients/auth/keyring.py

@@ -1,10 +1,9 @@
-import logging
 import typing
 from dataclasses import dataclass
 
 import keyring as _keyring
 from keyring.errors import NoKeyringError
-
+from flytekit import logger
 
 @dataclass
 class Credentials(object):
@@ -40,7 +39,7 @@ def store(credentials: Credentials) -> Credentials:
                 credentials.access_token,
             )
         except NoKeyringError as e:
-            logging.debug(f"KeyRing not available, tokens will not be cached. Error: {e}")
+            logger.debug(f"KeyRing not available, tokens will not be cached. Error: {e}")
         return credentials
 
     @staticmethod
@@ -49,7 +48,7 @@ def retrieve(for_endpoint: str) -> typing.Optional[Credentials]:
             refresh_token = _keyring.get_password(for_endpoint, KeyringStore._refresh_token_key)
             access_token = _keyring.get_password(for_endpoint, KeyringStore._access_token_key)
         except NoKeyringError as e:
-            logging.debug(f"KeyRing not available, tokens will not be cached. Error: {e}")
+            logger.debug(f"KeyRing not available, tokens will not be cached. Error: {e}")
             return None
 
         if not access_token:
@@ -62,4 +61,22 @@ def delete(for_endpoint: str):
             _keyring.delete_password(for_endpoint, KeyringStore._access_token_key)
             _keyring.delete_password(for_endpoint, KeyringStore._refresh_token_key)
         except NoKeyringError as e:
-            logging.debug(f"KeyRing not available, tokens will not be cached. Error: {e}")
+            logger.debug(f"KeyRing not available, tokens will not be cached. Error: {e}")
+
+
+# import flytekit
+# from flytekit.remote.remote import FlyteRemote
+# from flytekit.configuration import Config, PlatformConfig
+
+# cfg = Config(
+#     platform=PlatformConfig(
+#         endpoint="flyte.grid1.ard.grid.linkedin.com:31932",
+#         insecure=False,insecure_skip_verify=True,
+#         auth_mode="DeviceFlow",
+#         ca_cert_file_path="/etc/riddler/ca-bundle.crt",
+#     )
+# )
+
+# remote = FlyteRemote(config=cfg, default_project="openconnect-demo", default_domain="development")
+# client = remote.client
+# client.list_projects()

flytekit/clients/auth/token_client.py

@@ -1,6 +1,5 @@
 import base64
 import enum
-import logging
 import time
 import typing
 import urllib.parse
@@ -10,6 +9,7 @@
 import requests
 
 from flytekit.clients.auth.exceptions import AuthenticationError, AuthenticationPending
+from flytekit import logger
 
 utf_8 = "utf-8"
 
@@ -39,7 +39,7 @@ class DeviceCodeResponse:
     device_code: str
     user_code: str
     verification_uri: str
-    verification_uri_complete: str
+    # verification_uri_complete: str
     expires_in: int
     interval: int
 
@@ -49,7 +49,7 @@ def from_json_response(cls, j: typing.Dict) -> "DeviceCodeResponse":
             device_code=j["device_code"],
             user_code=j["user_code"],
             verification_uri=j["verification_uri"],
-            verification_uri_complete=j["verification_uri_complete"],
+            # verification_uri_complete=j["verification_uri_complete"],
             expires_in=j["expires_in"],
             interval=j["interval"],
         )
@@ -106,7 +106,7 @@ def get_token(
             err = j["error"]
             if err == error_auth_pending or err == error_slow_down:
                 raise AuthenticationPending(f"Token not yet available, try again in some time {err}")
-        logging.error("Status Code ({}) received from IDP: {}".format(response.status_code, response.text))
+        logger.error("Status Code ({}) received from IDP: {}".format(response.status_code, response.text))
         raise AuthenticationError("Status Code ({}) received from IDP: {}".format(response.status_code, response.text))
 
     j = response.json()

flytekit/clients/auth_helper.py

@@ -1,4 +1,3 @@
-import logging
 import ssl
 
 import grpc
@@ -18,7 +17,7 @@
 from flytekit.clients.grpc_utils.auth_interceptor import AuthUnaryInterceptor
 from flytekit.clients.grpc_utils.wrap_exception_interceptor import RetryExceptionWrapperInterceptor
 from flytekit.configuration import AuthType, PlatformConfig
-
+from flytekit import logger
 
 class RemoteClientConfigStore(ClientConfigStore):
     """
@@ -55,7 +54,7 @@ def get_authenticator(cfg: PlatformConfig, cfg_store: ClientConfigStore) -> Auth
         try:
             cfg_auth = AuthType[cfg_auth.upper()]
         except KeyError:
-            logging.warning(f"Authentication type {cfg_auth} does not exist, defaulting to standard")
+            logger.warning(f"Authentication type {cfg_auth} does not exist, defaulting to standard")
             cfg_auth = AuthType.STANDARD
 
     if cfg_auth == AuthType.STANDARD or cfg_auth == AuthType.PKCE:

flytekit/loggers.py

@@ -57,6 +57,7 @@
 flytekit_root_env_var = f"{LOGGING_ENV_VAR}_ROOT"
 level_from_env = os.getenv(LOGGING_ENV_VAR)
 root_level_from_env = os.getenv(flytekit_root_env_var)
+
 if root_level_from_env is not None:
     logger.setLevel(int(root_level_from_env))
 elif level_from_env is not None: