def test_redshift(self):
database = 'soda_test'
schema = 'public'
connection = psycopg2.connect(
user=os.getenv('SODA_REDSHIFT_USERNAME'),
password=os.getenv('SODA_REDSHIFT_PASSWORD'),
host=os.getenv('SODA_REDSHIFT_ENDPOINT'),
port='5439',
database=database)
try:
sql_update(connection, (
f'DROP TABLE IF EXISTS "{database}"."{schema}"."{self.table_name}"'
))
sql_update(connection, (
f'CREATE TABLE "{database}"."{schema}"."{self.table_name}" (\n'
f' "id" VARCHAR(255), \n'
f' "size" INTEGER );'))
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}"."{schema}"."{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}" \n'
f'WHERE "l" > 0'))
finally:
connection.close()
def test_postgres(self):
database = 'sodasql'
schema = 'public'
connection = psycopg2.connect(user='******',
host='localhost',
database=database,
options=f'-c search_path={schema}')
try:
sql_update(connection, (
f'DROP TABLE IF EXISTS "{database}"."{schema}"."{self.table_name}"'
))
sql_update(connection, (
f'CREATE TABLE "{database}"."{schema}"."{self.table_name}" (\n'
f' "id" VARCHAR(255), \n'
f' "size" INTEGER );'))
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}"."{schema}"."{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}" \n'
f'WHERE "l" > 0'))
finally:
connection.close()
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 test_athena(self):
Boto3Helper.filter_false_positive_boto3_warning()
aws_access_key_id = os.getenv('SODA_ATHENA_ACCESS_KEY_ID')
aws_secret_access_key = os.getenv('SODA_ATHENA_SECRET_ACCESS_KEY')
s3_staging_bucket = 'sodalite-test'
s3_staging_folder = 'soda_quotes_test'
database = 'sodalite_test'
schema = 'PUBLIC'
connection = pyathena.connect(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
s3_staging_dir=f's3://{s3_staging_bucket}/',
region_name='eu-west-1')
try:
sql_update(connection,
(f'CREATE DATABASE IF NOT EXISTS `{database}`'))
sql_update(
connection,
(f'DROP TABLE IF EXISTS `{database}`.`{self.table_name}`'))
sql_update(
connection,
(f'CREATE EXTERNAL TABLE `{database}`.`{self.table_name}` (\n'
f' `id` VARCHAR(255), \n'
f' `size` INTEGER ) \n'
f"LOCATION 's3://{s3_staging_bucket}/{s3_staging_folder}';"))
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}" \n'
f'WHERE "l" > 0'))
finally:
TestSqlQuotes.delete_athena_s3_staging_files(
aws_access_key_id, aws_secret_access_key, s3_staging_bucket,
s3_staging_folder)
connection.close()
def test_snowflake(self):
schema = 'PUBLIC'
connection = connector.connect(
user=os.getenv('SODA_SNOWFLAKE_USERNAME'),
password=os.getenv('SODA_SNOWFLAKE_PASSWORD'),
account='SODADATAPARTNER.eu-central-1',
warehouse='DEMO_WH',
schema=schema)
try:
database = 'soda_test_quotes_db'
sql_update(connection,
(f'CREATE DATABASE IF NOT EXISTS "{database}"'))
sql_update(connection, (
f'DROP TABLE IF EXISTS "{database}"."{schema}"."{self.table_name}"'
))
sql_update(connection, (
f'CREATE TABLE "{database}"."{schema}"."{self.table_name}" (\n'
f' "id" VARCHAR(255), \n'
f' "size" INTEGER );'))
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}"."{schema}"."{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}" \n'
f'WHERE "l" > 0'))
finally:
connection.close()
def sql_fetchone(self, sql) -> tuple:
return sql_fetchone(self.connection, sql)