diff --git a/sdk/python/kfp/_client.py b/sdk/python/kfp/_client.py index d2e9a111b848..3b60512cbd97 100644 --- a/sdk/python/kfp/_client.py +++ b/sdk/python/kfp/_client.py @@ -83,7 +83,7 @@ class Client(object): KUBE_PROXY_PATH = 'api/v1/namespaces/{}/services/ml-pipeline:http/proxy/' # TODO: Wrap the configurations for different authentication methods. - def __init__(self, host=None, client_id=None, namespace='kubeflow', other_client_id=None, other_client_secret=None): + def __init__(self, host=None, client_id=None, namespace='kubeflow', other_client_id=None, other_client_secret=None, existing_token=None): """Create a new instance of kfp client. Args: @@ -97,12 +97,14 @@ def __init__(self, host=None, client_id=None, namespace='kubeflow', other_client client_id: The client ID used by Identity-Aware Proxy. namespace: the namespace where the kubeflow pipeline system is run. other_client_id: The client ID used to obtain the auth codes and refresh tokens. - Reference: https://cloud.google.com/iap/docs/authentication-howto#authenticating_from_a_desktop_app. + Reference: https://cloud.google.com/iap/docs/authentication-howto#authenticating_from_a_desktop_app. other_client_secret: The client secret used to obtain the auth codes and refresh tokens. + existing_token: pass in token directly, it's used for cases better get token outside of SDK, e.x. GCP Cloud Functions + or caller already has a token """ host = host or os.environ.get(KF_PIPELINES_ENDPOINT_ENV) self._uihost = os.environ.get(KF_PIPELINES_UI_ENDPOINT_ENV, host) - config = self._load_config(host, client_id, namespace, other_client_id, other_client_secret) + config = self._load_config(host, client_id, namespace, other_client_id, other_client_secret, existing_token) api_client = kfp_server_api.api_client.ApiClient(config) _add_generated_apis(self, kfp_server_api, api_client) self._job_api = kfp_server_api.api.job_service_api.JobServiceApi(api_client) @@ -111,7 +113,7 @@ def __init__(self, host=None, client_id=None, namespace='kubeflow', other_client self._pipelines_api = kfp_server_api.api.pipeline_service_api.PipelineServiceApi(api_client) self._upload_api = kfp_server_api.api.PipelineUploadServiceApi(api_client) - def _load_config(self, host, client_id, namespace, other_client_id, other_client_secret): + def _load_config(self, host, client_id, namespace, other_client_id, other_client_secret, existing_token): config = kfp_server_api.configuration.Configuration() host = host or '' @@ -123,9 +125,26 @@ def _load_config(self, host, client_id, namespace, other_client_id, other_client token = None - # Obtain the tokens if it is IAP or inverse proxy. - # client_id is only used for IAP, so when the value is provided, we assume it's IAP. - if client_id: + # "existing_token" is designed to accept token generated outside of SDK. Here is an example. + # + # https://cloud.google.com/functions/docs/securing/function-identity + # https://cloud.google.com/endpoints/docs/grpc/service-account-authentication + # + # import requests + # import kfp + # + # def get_access_token(): + # url = 'http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/token' + # r = requests.get(url, headers={'Metadata-Flavor': 'Google'}) + # r.raise_for_status() + # access_token = r.json()['access_token'] + # return access_token + # + # client = kfp.Client(host='', existing_token=get_access_token()) + # + if existing_token: + token = existing_token + elif client_id: token = get_auth_token(client_id, other_client_id, other_client_secret) elif self._is_inverse_proxy_host(host): token = get_gcp_access_token()