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)