def test_bigquery(self):
database = 'sodalite'
account_info_json_str = os.getenv('BIGQUERY_ACCOUNT_INFO_JSON')
account_info_json_dict = json.loads(account_info_json_str)
credentials = Credentials.from_service_account_info(
account_info_json_dict)
project_id = account_info_json_dict['project_id']
client = bigquery.Client(project=project_id, credentials=credentials)
connection = dbapi.Connection(client)
try:
sql_update(
connection,
(f'DROP TABLE IF EXISTS `{database}`.`{self.table_name}`'))
sql_update(connection,
(f'CREATE TABLE `{database}`.`{self.table_name}` (\n'
f' `id` STRING, \n'
f' `size` INT64 );'))
sql_fetchone(
connection,
(f'WITH `{self.cte_table_name}` as ( \n'
f' SELECT "id" as `v`, "size" as `s`, LENGTH("id") as `l` \n'
f' FROM `{database}`.`{self.table_name}` \n'
f' WHERE `size` = 1 \n'
f' ORDER BY `size` ASC ) \n'
f'SELECT COUNT(DISTINCT("v")), COUNT("s") \n'
f'FROM `{self.cte_table_name}`;'))
finally:
connection.close()
def create_connection(self, *args, **kwargs):
credentials = Credentials.from_service_account_info(
self.account_info_dict)
project_id = self.account_info_dict['project_id']
self.client = bigquery.Client(project=project_id,
credentials=credentials)
return dbapi.Connection(self.client)
def create_connection(self):
try:
credentials = Credentials.from_service_account_info(
self.account_info_dict)
project_id = self.account_info_dict['project_id']
self.client = bigquery.Client(project=project_id,
credentials=credentials)
conn = dbapi.Connection(self.client)
return conn
except Exception as e:
self.try_to_raise_soda_sql_exception(e)
def __enter__(self) -> 'google.cloud.bigquery.dbapi.Connection':
from google.cloud.bigquery import dbapi
self.saved_credentials = os.environ.get(
BigQueryHook.credentials_env_var)
credentials_path = self.conn_params.extra.get('credentials_path')
if credentials_path:
os.environ[BigQueryHook.credentials_env_var] = credentials_path
self.conn = dbapi.Connection(self.client)
return self.conn
def create_connection(self):
try:
if self.__context_auth:
credentials = None
elif self.account_info_dict:
credentials = Credentials.from_service_account_info(self.account_info_dict, scopes=self.auth_scopes)
else:
raise Exception("Account_info_json or account_info_json_path or use_context_auth are not provided")
self.client = bigquery.Client(project=self.project_id, credentials=credentials)
conn = dbapi.Connection(self.client)
return conn
except Exception as e:
self.try_to_raise_soda_sql_exception(e)
def create_connection(self):
try:
auth_scopes = [
'https://www.googleapis.com/auth/bigquery',
'https://www.googleapis.com/auth/cloud-platform',
'https://www.googleapis.com/auth/drive'
]
credentials = Credentials.from_service_account_info(
self.account_info_dict, scopes=auth_scopes)
project_id = self.account_info_dict['project_id']
self.client = bigquery.Client(project=project_id,
credentials=credentials)
conn = dbapi.Connection(self.client)
return conn
except Exception as e:
self.try_to_raise_soda_sql_exception(e)
def get_db_conn():
global global_conn
if global_conn is None:
if database_target['type'] == "sqlserver":
db_pyodbc_connection_string = "Driver={SQL Server Native Client 11.0};Server=%s;Database=%s;Trusted_Connection=yes" % \
(database_target['host'], database_target['database'])
global_conn = pyodbc.connect(db_pyodbc_connection_string)
elif database_target['type'] == "sqlite":
schema_defs = database_target['schemas_and_paths']
db_sqlite_path = schema_defs['main']
global_conn = sqlite3.connect(db_sqlite_path)
elif database_target['type'] == "bigquery":
from google.cloud.bigquery import dbapi
client = get_bq_client()
global_conn = dbapi.Connection(client)
elif database_target['type'] == "snowflake":
global_conn = snowflake.connector.connect(
account=database_target['account'],
user=database_target['user'],
password=database_target['password'],
role=database_target['role'],
database=database_target['database'],
warehouse=database_target['warehouse'],
schema=database_target['schema'])
else:
raise "Unrecognized type: %s" % (database_target['type'], )
return global_conn
