def test_iam_role_partition(): """Tests the use of iam role with a custom partition""" aws_partition = 'aws-us-gov' aws_account_id = '000123456789' iam_role_name = 'redshiftrole' creds = 'aws_iam_role=arn:{0}:iam::{1}:role/{2}'.format( aws_partition, aws_account_id, iam_role_name, ) unload = dialect.UnloadFromSelect( select=sa.select([sa.func.count(table.c.id)]), unload_location='s3://bucket/key', aws_partition=aws_partition, aws_account_id=aws_account_id, iam_role_name=iam_role_name, ) expected_result = """ UNLOAD ('SELECT count(t1.id) AS count_1 FROM t1') TO 's3://bucket/key' CREDENTIALS '{creds}' """.format(creds=creds) assert clean(compile_query(unload)) == clean(expected_result)
def test_iam_role(stub_redshift_dialect): """Tests the use of iam role instead of access keys.""" aws_account_id = '000123456789' iam_role_name = 'redshiftrole' creds = 'aws_iam_role=arn:aws:iam::{0}:role/{1}'.format( aws_account_id, iam_role_name, ) unload = dialect.UnloadFromSelect( select=sa.select([sa.func.count(table.c.id)]), unload_location='s3://bucket/key', aws_account_id=aws_account_id, iam_role_name=iam_role_name, ) expected_result = """ UNLOAD ('SELECT count(t1.id) AS count_1 FROM t1') TO 's3://bucket/key' CREDENTIALS '{creds}' """.format(creds=creds) assert clean(compile_query(unload, stub_redshift_dialect)) == \ clean(expected_result)
def test_parquet_format__bad_options_crash(kwargs): """Verify we crash if we use the Parquet format with a bad option.""" unload = dialect.UnloadFromSelect(select=sa.select( [sa.func.count(table.c.id)]), unload_location='s3://bucket/key', access_key_id=access_key_id, secret_access_key=secret_access_key, format=dialect.Format.parquet, **kwargs) with pytest.raises(ValueError): compile_query(unload)
def test_iam_role_partition_validation(): """Tests the use of iam role with an invalid partition""" aws_partition = 'aws-invalid' aws_account_id = '000123456789' iam_role_name = 'redshiftrole' with pytest.raises(ValueError): dialect.UnloadFromSelect( select=sa.select([sa.func.count(table.c.id)]), unload_location='s3://bucket/key', aws_partition=aws_partition, aws_account_id=aws_account_id, iam_role_name=iam_role_name, )
def test_csv_format__bad_options_crash(delimiter, fixed_width): """Test that UnloadFromSelect crashes if you try to use DELIMITER and/or FIXEDWIDTH with the CSV format. """ unload = dialect.UnloadFromSelect(select=sa.select( [sa.func.count(table.c.id)]), unload_location='s3://bucket/key', access_key_id=access_key_id, secret_access_key=secret_access_key, format=dialect.Format.csv, delimiter=delimiter, fixed_width=fixed_width) with pytest.raises(ValueError): compile_query(unload)
def test_basic_unload_case(): """Tests that the simplest type of UnloadFromSelect works.""" unload = dialect.UnloadFromSelect( select=sa.select([sa.func.count(table.c.id)]), unload_location='s3://bucket/key', access_key_id=access_key_id, secret_access_key=secret_access_key, ) expected_result = """ UNLOAD ('SELECT count(t1.id) AS count_1 FROM t1') TO 's3://bucket/key' CREDENTIALS '{creds}' """.format(creds=creds) assert clean(compile_query(unload)) == clean(expected_result)
def test_parquet_format__basic(): """Basic successful test of unloading with the Parquet format.""" unload = dialect.UnloadFromSelect( select=sa.select([sa.func.count(table.c.id)]), unload_location='s3://bucket/key', access_key_id=access_key_id, secret_access_key=secret_access_key, format=dialect.Format.parquet, ) expected_result = """ UNLOAD ('SELECT count(t1.id) AS count_1 FROM t1') TO 's3://bucket/key' CREDENTIALS '{creds}' FORMAT AS PARQUET """.format(creds=creds) assert clean(compile_query(unload)) == clean(expected_result)
def test_csv_format__basic(): """Tests that UnloadFromSelect uses the format option correctly.""" unload = dialect.UnloadFromSelect(select=sa.select( [sa.func.count(table.c.id)]), unload_location='s3://bucket/key', access_key_id=access_key_id, secret_access_key=secret_access_key, format=dialect.Format.csv) expected_result = """ UNLOAD ('SELECT count(t1.id) AS count_1 FROM t1') TO 's3://bucket/key' CREDENTIALS '{creds}' FORMAT AS CSV """.format(creds=creds) assert clean(compile_query(unload)) == clean(expected_result)
def test_all_redshift_options(): """Tests that UnloadFromSelect handles all options correctly.""" unload = dialect.UnloadFromSelect( sa.select([sa.func.count(table.c.id)]), unload_location='s3://bucket/key', access_key_id=access_key_id, secret_access_key=secret_access_key, manifest=True, delimiter=',', fixed_width=[ ('count_1', 50), ], encrypted=True, gzip=True, add_quotes=True, null='---', escape=True, allow_overwrite=True, parallel=False, region='us-west-2', max_file_size=10 * 1024**2, ) expected_result = """ UNLOAD ('SELECT count(t1.id) AS count_1 FROM t1') TO 's3://bucket/key' CREDENTIALS '{creds}' MANIFEST DELIMITER AS ',' ENCRYPTED FIXEDWIDTH AS 'count_1:50' GZIP ADDQUOTES NULL AS '---' ESCAPE ALLOWOVERWRITE PARALLEL OFF REGION 'us-west-2' MAXFILESIZE 10.0 MB """.format(creds=creds) assert clean(compile_query(unload)) == clean(expected_result)
def test_iam_role_arns_single(): """Tests the use of a single iam role arn instead of access keys.""" iam_role_arns = 'arn:aws:iam::000123456789:role/redshiftrole' creds = 'aws_iam_role=arn:aws:iam::000123456789:role/redshiftrole' unload = dialect.UnloadFromSelect( select=sa.select([sa.func.count(table.c.id)]), unload_location='s3://bucket/key', iam_role_arns=iam_role_arns, ) expected_result = """ UNLOAD ('SELECT count(t1.id) AS count_1 FROM t1') TO 's3://bucket/key' CREDENTIALS '{creds}' """.format(creds=creds) assert clean(compile_query(unload)) == clean(expected_result)
def test_all_redshift_options_with_header(): """Tests that UnloadFromSelect handles all options correctly.""" unload = dialect.UnloadFromSelect(sa.select([sa.func.count(table.c.id)]), unload_location='s3://bucket/key', access_key_id=access_key_id, secret_access_key=secret_access_key, manifest=True, header=True, delimiter=',', encrypted=True, gzip=True, add_quotes=True, null='---', escape=True, allow_overwrite=True, parallel=False, region='ap-northeast-2') expected_result = """ UNLOAD ('SELECT count(t1.id) AS count_1 FROM t1') TO 's3://bucket/key' CREDENTIALS '{creds}' MANIFEST HEADER DELIMITER AS ',' ENCRYPTED GZIP ADDQUOTES NULL AS '---' ESCAPE ALLOWOVERWRITE PARALLEL OFF REGION 'ap-northeast-2' """.format(creds=creds) assert clean(compile_query(unload)) == clean(expected_result)