Exemplo n.º 1
0
    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()
Exemplo n.º 2
0
 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)
Exemplo n.º 3
0
 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)
Exemplo n.º 4
0
    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
Exemplo n.º 5
0
    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)
Exemplo n.º 6
0
 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)
Exemplo n.º 7
0
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