| 133 | return "KAGGLE_DATA_PROXY_TOKEN" in os.environ |
| 134 | |
| 135 | def init_bigquery(): |
| 136 | from google.cloud import bigquery |
| 137 | |
| 138 | if not (is_proxy_token_set() or is_user_secrets_token_set()): |
| 139 | return bigquery |
| 140 | |
| 141 | # If this Notebook has bigquery integration on startup, preload the Kaggle Credentials |
| 142 | # object for magics to work. |
| 143 | if get_integrations().has_bigquery(): |
| 144 | from google.cloud.bigquery import magics |
| 145 | magics.context.credentials = KaggleKernelCredentials() |
| 146 | |
| 147 | def monkeypatch_bq(bq_client, *args, **kwargs): |
| 148 | from kaggle_gcp import get_integrations, PublicBigqueryClient, KaggleKernelCredentials |
| 149 | specified_credentials = kwargs.get('credentials') |
| 150 | has_bigquery = get_integrations().has_bigquery() |
| 151 | # Prioritize passed in project id, but if it is missing look for env var. |
| 152 | arg_project = kwargs.get('project') |
| 153 | explicit_project_id = arg_project or os.environ.get(environment_vars.PROJECT) |
| 154 | # This is a hack to get around the bug in google-cloud library. |
| 155 | # Remove these two lines once this is resolved: |
| 156 | # https://github.com/googleapis/google-cloud-python/issues/8108 |
| 157 | if explicit_project_id: |
| 158 | logging.info(f"Explicit project set to {explicit_project_id}") |
| 159 | kwargs['project'] = explicit_project_id |
| 160 | if explicit_project_id is None and specified_credentials is None and not has_bigquery: |
| 161 | msg = "Using Kaggle's public dataset BigQuery integration." |
| 162 | logging.info(msg) |
| 163 | print(msg) |
| 164 | return PublicBigqueryClient(*args, **kwargs) |
| 165 | else: |
| 166 | if specified_credentials is None: |
| 167 | logging.info("No credentials specified, using KaggleKernelCredentials.") |
| 168 | kwargs['credentials'] = KaggleKernelCredentials() |
| 169 | if (not has_bigquery): |
| 170 | logging.info("No bigquery integration found, creating client anyways.") |
| 171 | print('Please ensure you have selected a BigQuery ' |
| 172 | 'account in the Notebook Add-ons menu.') |
| 173 | if explicit_project_id is None: |
| 174 | logging.info("No project specified while using the unmodified client.") |
| 175 | print('Please ensure you specify a project id when creating the client' |
| 176 | ' in order to use your BigQuery account.') |
| 177 | kwargs['client_info'] = set_kaggle_user_agent(kwargs.get('client_info')) |
| 178 | return bq_client(*args, **kwargs) |
| 179 | |
| 180 | # Monkey patches BigQuery client creation to use proxy or user-connected GCP account. |
| 181 | # Deprecated in favor of Kaggle.DataProxyClient(). |
| 182 | # TODO: Remove this once uses have migrated to that new interface. |
| 183 | bq_client = bigquery.Client |
| 184 | if (not has_been_monkeypatched(bigquery.Client)): |
| 185 | bigquery.Client = lambda *args, **kwargs: monkeypatch_bq( |
| 186 | bq_client, *args, **kwargs) |
| 187 | return bigquery |
| 188 | |
| 189 | # Monkey patch for aiplatform init |
| 190 | # eg |