docker-python/patches/sitecustomize.py at master · HadokenCode/docker-python · GitHub

27 lines (22 loc) · 1 KB

# Monkey patches BigQuery client creation to use proxy.

# Import tensorflow and torch before anything else. This is a hacky workaround to an error on dlopen

# reporting a limit on static TLS, tracked in:

# tensorflow: https://github.com/tensorflow/tensorflow/issues/19010

# torch: https://github.com/pytorch/pytorch/issues/2575

import tensorflow

import torch

import os

kaggle_proxy_token = os.getenv("KAGGLE_DATA_PROXY_TOKEN")

if kaggle_proxy_token:

from google.auth import credentials

from google.cloud import bigquery

from google.cloud.bigquery._http import Connection

Connection.API_BASE_URL = os.getenv("KAGGLE_DATA_PROXY_URL")

Connection._EXTRA_HEADERS["X-KAGGLE-PROXY-DATA"] = kaggle_proxy_token

bq_client = bigquery.Client

bigquery.Client = lambda *args, **kwargs: bq_client(

*args,

credentials=credentials.AnonymousCredentials(),

project=os.getenv("KAGGLE_DATA_PROXY_PROJECT"),

**kwargs)

credentials.AnonymousCredentials.refresh = lambda *args: None

Expand file tree