def init_automl(): is_user_secrets_token_set = "KAGGLE_USER_SECRETS_TOKEN" in os.environ from google.cloud import automl_v1beta1 as automl if not is_user_secrets_token_set: return automl from kaggle_gcp import get_integrations if not get_integrations().has_automl(): return automl from kaggle_secrets import GcpTarget from kaggle_gcp import KaggleKernelCredentials kaggle_kernel_credentials = KaggleKernelCredentials( target=GcpTarget.AUTOML) # The AutoML client library exposes 4 different client classes (AutoMlClient, # TablesClient, PredictionServiceClient and GcsClient), so patch each of them. # The same KaggleKernelCredentials are passed to all of them. monkeypatch_client(automl.AutoMlClient, kaggle_kernel_credentials) monkeypatch_client(automl.TablesClient, kaggle_kernel_credentials) monkeypatch_client(automl.PredictionServiceClient, kaggle_kernel_credentials) # TODO(markcollins): The GcsClient in the AutoML client library version # 0.5.0 doesn't handle credentials properly. I wrote PR: # https://github.com/googleapis/google-cloud-python/pull/9299 # to address this issue. Add patching for GcsClient when we get a version of # the library that includes the fixes. return automl
def test_project_with_connected_account(self, mock_access_token, ApiUrlMock): self._setup_mocks(ApiUrlMock) env = EnvironmentVarGuard() env.set('KAGGLE_USER_SECRETS_TOKEN', 'foobar') with env: client = bigquery.Client( project='ANOTHER_PROJECT', credentials=KaggleKernelCredentials()) self._test_integration(client)
def monkeypatch_gcs(self, *args, **kwargs): specified_credentials = kwargs.get('credentials') if specified_credentials is None: Log.info( "No credentials specified, using KaggleKernelCredentials.") kwargs['credentials'] = KaggleKernelCredentials( target=GcpTarget.GCS) return gcs_client_init(self, *args, **kwargs)
def test_project_with_connected_account(self, mock_access_token): env = EnvironmentVarGuard() env.set('KAGGLE_USER_SECRETS_TOKEN', 'foobar') with env: client = bigquery.Client( project='ANOTHER_PROJECT', credentials=KaggleKernelCredentials(), client_options={"api_endpoint": TestBigQuery.API_BASE_URL}) self._test_integration(client)
def init_bigquery(): from google.auth import credentials, environment_vars from google.cloud import bigquery from google.cloud.bigquery._http import Connection # TODO: Update this to the correct kaggle.gcp path once we no longer inject modules # from the worker. from kaggle_gcp import get_integrations, PublicBigqueryClient, KaggleKernelCredentials # If this Kernel has bigquery integration on startup, preload the Kaggle Credentials # object for magics to work. if get_integrations().has_bigquery(): from google.cloud.bigquery import magics magics.context.credentials = KaggleKernelCredentials() def monkeypatch_bq(bq_client, *args, **kwargs): specified_credentials = kwargs.get('credentials') has_bigquery = get_integrations().has_bigquery() # Prioritize passed in project id, but if it is missing look for env var. arg_project = kwargs.get('project') explicit_project_id = arg_project or os.environ.get( environment_vars.PROJECT) # This is a hack to get around the bug in google-cloud library. # Remove these two lines once this is resolved: # https://github.com/googleapis/google-cloud-python/issues/8108 if explicit_project_id: Log.info(f"Explicit project set to {explicit_project_id}") kwargs['project'] = explicit_project_id if explicit_project_id is None and specified_credentials is None and not has_bigquery: msg = "Using Kaggle's public dataset BigQuery integration." Log.info(msg) print(msg) return PublicBigqueryClient(*args, **kwargs) else: if specified_credentials is None: Log.info( "No credentials specified, using KaggleKernelCredentials.") kwargs['credentials'] = KaggleKernelCredentials() if (not has_bigquery): Log.info( "No bigquery integration found, creating client anyways." ) print('Please ensure you have selected a BigQuery ' 'account in the Kernels Settings sidebar.') return bq_client(*args, **kwargs) # Monkey patches BigQuery client creation to use proxy or user-connected GCP account. # Deprecated in favor of Kaggle.DataProxyClient(). # TODO: Remove this once uses have migrated to that new interface. bq_client = bigquery.Client bigquery.Client = lambda *args, **kwargs: monkeypatch_bq( bq_client, *args, **kwargs)
def test_simultaneous_clients(self, mock_access_token, ApiUrlMock): self._setup_mocks(ApiUrlMock) env = EnvironmentVarGuard() env.set('KAGGLE_USER_SECRETS_TOKEN', 'foobar') with env: proxy_client = bigquery.Client() bq_client = bigquery.Client( project='ANOTHER_PROJECT', credentials=KaggleKernelCredentials()) self._test_integration(bq_client) # Verify that proxy client is still going to proxy to ensure global Connection # isn't being modified. self.assertNotEqual(type(proxy_client._connection), KaggleKernelCredentials) self.assertEqual(type(proxy_client._connection), _DataProxyConnection)
def init_translation_v3(): # Translate v3 exposes different client than translate v2. from google.cloud import translate_v3 if not is_user_secrets_token_set(): return translate_v3 from kaggle_gcp import get_integrations if not get_integrations().has_cloudai(): return translate_v3 from kaggle_secrets import GcpTarget kernel_credentials = KaggleKernelCredentials(target=GcpTarget.CLOUDAI) monkeypatch_client(translate_v3.TranslationServiceClient, kernel_credentials) return translate_v3
def monkeypatch_bq(bq_client, *args, **kwargs): data_proxy_project = os.getenv("KAGGLE_DATA_PROXY_PROJECT") specified_project = kwargs.get('project') specified_credentials = kwargs.get('credentials') kernel_integrations = get_integrations() if specified_project is None and specified_credentials is None and not kernel_integrations.has_bigquery(): print("Using Kaggle's public dataset BigQuery integration.") return PublicBigqueryClient(*args, **kwargs) else: if specified_credentials is None: kwargs['credentials'] = KaggleKernelCredentials() return bq_client(*args, **kwargs)
def monkeypatch_bq(bq_client, *args, **kwargs): specified_project = kwargs.get('project') specified_credentials = kwargs.get('credentials') has_bigquery = get_integrations().has_bigquery() if specified_project is None and specified_credentials is None and not has_bigquery: print("Using Kaggle's public dataset BigQuery integration.") return PublicBigqueryClient(*args, **kwargs) else: if specified_credentials is None: kwargs['credentials'] = KaggleKernelCredentials() if (not has_bigquery): print('Please ensure you have selected a BigQuery ' 'account in the Kernels Settings sidebar.') return bq_client(*args, **kwargs)
def init_gcs(): is_user_secrets_token_set = "KAGGLE_USER_SECRETS_TOKEN" in os.environ from google.cloud import storage if not is_user_secrets_token_set: return storage from kaggle_gcp import get_integrations if not get_integrations().has_gcs(): return storage from kaggle_secrets import GcpTarget from kaggle_gcp import KaggleKernelCredentials monkeypatch_client(storage.Client, KaggleKernelCredentials(target=GcpTarget.GCS)) return storage
def init_ucaip(): from google.cloud import aiplatform if not is_user_secrets_token_set(): return from kaggle_gcp import get_integrations if not get_integrations().has_cloudai(): return from kaggle_secrets import GcpTarget from kaggle_gcp import KaggleKernelCredentials kaggle_kernel_credentials = KaggleKernelCredentials(target=GcpTarget.CLOUDAI) # Patch the ucaip init method, this flows down to all ucaip services monkeypatch_aiplatform_init(aiplatform, kaggle_kernel_credentials)
def test_simultaneous_clients(self, mock_access_token): env = EnvironmentVarGuard() env.set('KAGGLE_USER_SECRETS_TOKEN', 'foobar') with env: proxy_client = bigquery.Client( client_options={"api_endpoint": TestBigQuery.API_BASE_URL}) bq_client = bigquery.Client( project='ANOTHER_PROJECT', credentials=KaggleKernelCredentials(), client_options={"api_endpoint": TestBigQuery.API_BASE_URL}) self._test_integration(bq_client) # Verify that proxy client is still going to proxy to ensure global Connection # isn't being modified. self.assertNotEqual(type(proxy_client._connection), KaggleKernelCredentials) self.assertEqual(type(proxy_client._connection), _DataProxyConnection)
def test_default_target(self): creds = KaggleKernelCredentials() self.assertEqual(GcpTarget.BIGQUERY, creds.target)