def to_redshift(self, table_name, username=None, password=None, host=None, db=None, port=None, **copy_args): """ Write a table to a Redshift database. Note, this requires you to pass AWS S3 credentials or store them as environmental variables. Args: table_name: str The table name and schema (``my_schema.my_table``) to point the file. username: str Required if env variable ``REDSHIFT_USERNAME`` not populated password: str Required if env variable ``REDSHIFT_PASSWORD`` not populated host: str Required if env variable ``REDSHIFT_HOST`` not populated db: str Required if env variable ``REDSHIFT_DB`` not populated port: int Required if env variable ``REDSHIFT_PORT`` not populated. Port 5439 is typical. \**copy_args: kwargs See :func:`~parsons.databases.Redshift.copy`` for options. Returns: ``None`` """ # noqa: W605 from parsons.databases.redshift import Redshift rs = Redshift(username=username, password=password, host=host, db=db, port=port) rs.copy(self, table_name, **copy_args)
class TestRedshiftDB(unittest.TestCase): def setUp(self): self.temp_schema = TEMP_SCHEMA self.rs = Redshift() self.tbl = Table([['ID', 'Name'], [1, 'Jim'], [2, 'John'], [3, 'Sarah']]) # Create a schema, create a table, create a view setup_sql = f""" drop schema if exists {self.temp_schema} cascade; create schema {self.temp_schema}; """ other_sql = f""" create table {self.temp_schema}.test (id int,name varchar(5)); create view {self.temp_schema}.test_view as ( select * from {self.temp_schema}.test ); """ self.rs.query(setup_sql) self.rs.query(other_sql) self.s3 = S3() self.temp_s3_bucket = os.environ['S3_TEMP_BUCKET'] self.temp_s3_prefix = 'test/' def tearDown(self): # Drop the view, the table and the schema teardown_sql = f""" drop schema if exists {self.temp_schema} cascade; """ self.rs.query(teardown_sql) # Remove all test objects from S3 for key in self.s3.list_keys(self.temp_s3_bucket, self.temp_s3_prefix): self.s3.remove_file(self.temp_s3_bucket, key) def test_query(self): # Check that query sending back expected result r = self.rs.query('select 1') self.assertEqual(r[0]['?column?'], 1) def test_query_with_parameters(self): table_name = f"{self.temp_schema}.test" self.tbl.to_redshift(table_name, if_exists='append') sql = f"select * from {table_name} where name = %s" name = 'Sarah' r = self.rs.query(sql, parameters=[name]) self.assertEqual(r[0]['name'], name) sql = f"select * from {table_name} where name in (%s, %s)" names = ['Sarah', 'John'] r = self.rs.query(sql, parameters=names) self.assertEqual(r.num_rows, 2) def test_schema_exists(self): self.assertTrue(self.rs.schema_exists(self.temp_schema)) self.assertFalse(self.rs.schema_exists('nonsense')) def test_table_exists(self): # Check if table_exists finds a table that exists self.assertTrue(self.rs.table_exists(f'{self.temp_schema}.test')) # Check if table_exists is case insensitive self.assertTrue(self.rs.table_exists(f'{self.temp_schema.upper()}.TEST')) # Check if table_exists doesn't find a table that doesn't exists self.assertFalse(self.rs.table_exists(f'{self.temp_schema}.test_fake')) # Check if table_exists finds a table that exists self.assertTrue(self.rs.table_exists(f'{self.temp_schema}.test_view')) # Check if table_exists doesn't find a view that doesn't exists self.assertFalse(self.rs.table_exists(f'{self.temp_schema}.test_view_fake')) # Check that the view kwarg works self.assertFalse(self.rs.table_exists(f'{self.temp_schema}.test_view', view=False)) def test_temp_s3_create(self): key = self.rs.temp_s3_copy(self.tbl) # Test that you can get the object self.s3.get_file(self.temp_s3_bucket, key) # Try to delete the object self.rs.temp_s3_delete(key) def test_copy_s3(self): # To Do pass def test_copy(self): # Copy a table self.rs.copy(self.tbl, f'{self.temp_schema}.test_copy', if_exists='drop') # Test that file exists r = self.rs.query(f"select * from {self.temp_schema}.test_copy where name='Jim'") self.assertEqual(r[0]['id'], 1) # Copy to the same table, to verify that the "truncate" flag works. self.rs.copy(self.tbl, f'{self.temp_schema}.test_copy', if_exists='truncate') rows = self.rs.query(f"select count(*) from {self.temp_schema}.test_copy") self.assertEqual(rows[0]['count'], 3) # Copy to the same table, to verify that the "drop" flag works. self.rs.copy(self.tbl, f'{self.temp_schema}.test_copy', if_exists='drop') # Verify that a warning message prints when a DIST/SORT key is omitted with LogCapture() as lc: self.rs.copy( self.tbl, f'{self.temp_schema}.test_copy', if_exists='drop', sortkey='Name') desired_log = [log for log in lc.records if "optimize your queries" in log.msg][0] self.assertTrue("DIST" in desired_log.msg) self.assertFalse("SORT" in desired_log.msg) def test_upsert(self): # Create a target table when no target table exists self.rs.upsert(self.tbl, f'{self.temp_schema}.test_copy', 'ID') # Run upsert upsert_tbl = Table([['id', 'name'], [1, 'Jane'], [5, 'Bob']]) self.rs.upsert(upsert_tbl, f'{self.temp_schema}.test_copy', 'ID') # Make sure that it is the expected table expected_tbl = Table([['id', 'name'], [1, 'Jane'], [2, 'John'], [3, 'Sarah'], [5, 'Bob']]) updated_tbl = self.rs.query(f'select * from {self.temp_schema}.test_copy order by id;') assert_matching_tables(expected_tbl, updated_tbl) # Try to run it with a bad primary key self.rs.query(f"INSERT INTO {self.temp_schema}.test_copy VALUES (1, 'Jim')") self.assertRaises( ValueError, self.rs.upsert, upsert_tbl, f'{self.temp_schema}.test_copy', 'ID') # Now try and upsert using two primary keys upsert_tbl = Table([['id', 'name'], [1, 'Jane']]) self.rs.upsert(upsert_tbl, f'{self.temp_schema}.test_copy', ['id', 'name']) # Make sure our table looks like we expect expected_tbl = Table([['id', 'name'], [2, 'John'], [3, 'Sarah'], [5, 'Bob'], [1, 'Jim'], [1, 'Jane']]) updated_tbl = self.rs.query(f'select * from {self.temp_schema}.test_copy order by id;') assert_matching_tables(expected_tbl, updated_tbl) # Try to run it with a bad primary key self.rs.query(f"INSERT INTO {self.temp_schema}.test_copy VALUES (1, 'Jim')") self.assertRaises(ValueError, self.rs.upsert, upsert_tbl, f'{self.temp_schema}.test_copy', ['ID', 'name']) self.rs.query(f'truncate table {self.temp_schema}.test_copy') # Run upsert with nonmatching datatypes upsert_tbl = Table([['id', 'name'], [3, 600], [6, 9999]]) self.rs.upsert(upsert_tbl, f'{self.temp_schema}.test_copy', 'ID') # Make sure our table looks like we expect expected_tbl = Table([['id', 'name'], [3, '600'], [6, '9999']]) updated_tbl = self.rs.query(f'select * from {self.temp_schema}.test_copy order by id;') assert_matching_tables(expected_tbl, updated_tbl) # Run upsert requiring column resize upsert_tbl = Table([['id', 'name'], [7, 'this name is very long']]) self.rs.upsert(upsert_tbl, f'{self.temp_schema}.test_copy', 'ID') # Make sure our table looks like we expect expected_tbl = Table([['id', 'name'], [3, '600'], [6, '9999'], [7, 'this name is very long']]) updated_tbl = self.rs.query(f'select * from {self.temp_schema}.test_copy order by id;') assert_matching_tables(expected_tbl, updated_tbl) def test_unload(self): # Copy a table to Redshift self.rs.copy(self.tbl, f'{self.temp_schema}.test_copy', if_exists='drop') # Unload a table to S3 self.rs.unload(f'select * from {self.temp_schema}.test_copy', self.temp_s3_bucket, 'unload_test') # Check that files are there self.assertTrue(self.s3.key_exists(self.temp_s3_bucket, 'unload_test')) def test_to_from_redshift(self): # Test the parsons table methods table_name = f'{self.temp_schema}.test_copy' self.tbl.to_redshift(table_name, if_exists='drop') sql = f"SELECT * FROM {table_name} ORDER BY id" result_tbl = Table.from_redshift(sql) # Don't bother checking columns names, since those were tweaked en route to Redshift. assert_matching_tables(self.tbl, result_tbl, ignore_headers=True) def test_generate_manifest(self): # Add some tables to buckets self.tbl.to_s3_csv(self.temp_s3_bucket, f'{self.temp_s3_prefix}test_file_01.csv') self.tbl.to_s3_csv(self.temp_s3_bucket, f'{self.temp_s3_prefix}test_file_02.csv') self.tbl.to_s3_csv(self.temp_s3_bucket, f'{self.temp_s3_prefix}test_file_03.csv') self.tbl.to_s3_csv(self.temp_s3_bucket, f'{self.temp_s3_prefix}dont_include.csv') # Copy in a table to generate the headers and table self.rs.copy(self.tbl, f'{self.temp_schema}.test_copy', if_exists='drop') # Generate the manifest manifest_key = f'{self.temp_s3_prefix}test_manifest.json' manifest = self.rs.generate_manifest(self.temp_s3_bucket, prefix=f'{self.temp_s3_prefix}test_file', manifest_bucket=self.temp_s3_bucket, manifest_key=manifest_key) # Validate path formatted correctly valid_url = f's3://{self.temp_s3_bucket}/{self.temp_s3_prefix}test_file_01.csv' self.assertEqual(manifest['entries'][0]['url'], valid_url) # Validate that there are three files self.assertEqual(len(manifest['entries']), 3) # Validate that manifest saved to bucket keys = self.s3.list_keys(self.temp_s3_bucket, prefix=f'{self.temp_s3_prefix}test_manifest') self.assertTrue(manifest_key in keys) def test_move_table(self): # Run the method and check that new table created self.rs.move_table(f'{self.temp_schema}.test', f'{self.temp_schema}.test2') self.assertTrue(self.rs.table_exists(f'{self.temp_schema}.test2')) # Run the method again, but drop original self.rs.move_table( f'{self.temp_schema}.test2', f'{self.temp_schema}.test3', drop_source_table=True) self.assertFalse(self.rs.table_exists(f'{self.temp_schema}.test2')) def test_get_tables(self): tbls_list = self.rs.get_tables(schema=self.temp_schema) exp = ['schemaname', 'tablename', 'tableowner', 'tablespace', 'hasindexes', 'hasrules', 'hastriggers'] self.assertTrue(validate_list(exp, tbls_list)) def test_get_table_stats(self): tbls_list = self.rs.get_table_stats(schema=self.temp_schema) exp = [ 'database', 'schema', 'table_id', 'table', 'encoded', 'diststyle', 'sortkey1', 'max_varchar', 'sortkey1_enc', 'sortkey_num', 'size', 'pct_used', 'empty', 'unsorted', 'stats_off', 'tbl_rows', 'skew_sortkey1', 'skew_rows', 'estimated_visible_rows', 'risk_event', 'vacuum_sort_benefit'] # Having some issues testing that the filter is working correctly, as it # takes a little bit of time for a table to show in this table and is beating # the test suite. I feel confident that it works though. self.assertTrue(validate_list(exp, tbls_list)) def test_get_views(self): # Assert that get_views returns filtered views # Assert that it works with schema filter views = self.rs.get_views(schema=self.temp_schema) expected_row = (self.temp_schema, 'test_view', f'SELECT test.id, test.name FROM {self.temp_schema}.test;') self.assertEqual(views.data[0], expected_row) def test_get_queries(self): # Validate that columns match expected columns queries_list = self.rs.get_queries() exp = ['user', 'pid', 'xid', 'query', 'service_class', 'slot', 'start', 'state', 'queue_sec', 'exec_sec', 'cpu_sec', 'read_mb', 'spill_mb', 'return_rows', 'nl_rows', 'sql', 'alert'] self.assertTrue(validate_list(exp, queries_list)) def test_get_row_count(self): table_name = f'{self.temp_schema}.test_row_count' self.rs.copy(self.tbl, table_name, if_exists='drop') count = self.rs.get_row_count(table_name) self.assertEqual(count, 3) def test_rename_table(self): self.rs.rename_table(self.temp_schema + '.test', 'test2') # Test that renamed table exists self.assertTrue(self.rs.table_exists(self.temp_schema + '.test2')) # Test that old table name does not exist self.assertFalse(self.rs.table_exists(self.temp_schema + '.test')) def test_union_tables(self): # Copy in two tables self.rs.copy(self.tbl, f'{self.temp_schema}.union_base1', if_exists='drop') self.rs.copy(self.tbl, f'{self.temp_schema}.union_base2', if_exists='drop') # Union all the two tables and check row count self.rs.union_tables(f'{self.temp_schema}.union_all', [f'{self.temp_schema}.union_base1', f'{self.temp_schema}.union_base2']) self.assertEqual(self.rs.query(f"select * from {self.temp_schema}.union_all").num_rows, 6) # Union the two tables and check row count self.rs.union_tables(f'{self.temp_schema}.union_test', [f'{self.temp_schema}.union_base1', f'{self.temp_schema}.union_base2'], union_all=False) self.assertEqual(self.rs.query(f"select * from {self.temp_schema}.union_test").num_rows, 3) def test_populate_table_from_query(self): # Populate the source table source_table = f'{self.temp_schema}.test_source' self.rs.copy(self.tbl, source_table, if_exists='drop') query = f"SELECT * FROM {source_table}" # Populate the table dest_table = f'{self.temp_schema}.test_dest' self.rs.populate_table_from_query(query, dest_table) # Verify rows = self.rs.query(f"select count(*) from {dest_table}") self.assertEqual(rows[0]['count'], 3) # Try with if_exists='truncate' self.rs.populate_table_from_query(query, dest_table, if_exists='truncate') rows = self.rs.query(f"select count(*) from {dest_table}") self.assertEqual(rows[0]['count'], 3) # Try with if_exists='drop', and a distkey self.rs.populate_table_from_query(query, dest_table, if_exists='drop', distkey='id') rows = self.rs.query(f"select count(*) from {dest_table}") self.assertEqual(rows[0]['count'], 3) # Try with if_exists='fail' self.assertRaises( ValueError, self.rs.populate_table_from_query, query, dest_table, if_exists='fail') def test_duplicate_table(self): # Populate the source table source_table = f'{self.temp_schema}.test_source' self.rs.copy(self.tbl, source_table, if_exists='drop') # Duplicate the table dest_table = f'{self.temp_schema}.test_dest' self.rs.duplicate_table(source_table, dest_table) # Verify rows = self.rs.query(f"select count(*) from {dest_table}") self.assertEqual(rows[0]['count'], 3) # Try with if_exists='truncate' self.rs.duplicate_table(source_table, dest_table, if_exists='truncate') rows = self.rs.query(f"select count(*) from {dest_table}") self.assertEqual(rows[0]['count'], 3) # Try with if_exists='drop' self.rs.duplicate_table(source_table, dest_table, if_exists='drop') rows = self.rs.query(f"select count(*) from {dest_table}") self.assertEqual(rows[0]['count'], 3) # Try with if_exists='append' self.rs.duplicate_table(source_table, dest_table, if_exists='append') rows = self.rs.query(f"select count(*) from {dest_table}") self.assertEqual(rows[0]['count'], 6) # Try with if_exists='fail' self.assertRaises(ValueError, self.rs.duplicate_table, source_table, dest_table, if_exists='fail') # Try with invalid if_exists arg self.assertRaises(ValueError, self.rs.duplicate_table, source_table, dest_table, if_exists='nonsense') def test_get_max_value(self): date_tbl = Table([['id', 'date_modified'], [1, '2020-01-01'], [2, '1900-01-01']]) self.rs.copy(date_tbl, f'{self.temp_schema}.test_date') # Test return string self.assertEqual(self.rs.get_max_value(f'{self.temp_schema}.test_date', 'date_modified'), '2020-01-01') def test_get_columns(self): cols = self.rs.get_columns(self.temp_schema, 'test') # id int,name varchar(5) expected_cols = { 'id': { 'data_type': 'int', 'max_length': None, 'max_precision': 32, 'max_scale': 0, 'is_nullable': True}, 'name': { 'data_type': 'character varying', 'max_length': 5, 'max_precision': None, 'max_scale': None, 'is_nullable': True}, } self.assertEqual(cols, expected_cols) def test_get_object_type(self): # Test a table expected_type_table = "table" actual_type_table = self.rs.get_object_type("pg_catalog.pg_class") self.assertEqual(expected_type_table, actual_type_table) # Test a view expected_type_view = "view" actual_type_view = self.rs.get_object_type("pg_catalog.pg_views") self.assertEqual(expected_type_view, actual_type_view) # Test a nonexisting table expected_type_fake = None actual_type_fake = self.rs.get_object_type("someschema.faketable") self.assertEqual(expected_type_fake, actual_type_fake) def test_is_view(self): self.assertTrue(self.rs.is_view("pg_catalog.pg_views")) self.assertFalse(self.rs.is_view("pg_catalog.pg_class")) def test_is_table(self): self.assertTrue(self.rs.is_table("pg_catalog.pg_class")) self.assertFalse(self.rs.is_table("pg_catalog.pg_views")) def test_get_table_definition(self): expected_table_def = ( '--DROP TABLE pg_catalog.pg_amop;' '\nCREATE TABLE IF NOT EXISTS pg_catalog.pg_amop' '\n(' '\n\tamopclaid OID NOT NULL ENCODE RAW' '\n\t,amopsubtype OID NOT NULL ENCODE RAW' '\n\t,amopstrategy SMALLINT NOT NULL ENCODE RAW' '\n\t,amopreqcheck BOOLEAN NOT NULL ENCODE RAW' '\n\t,amopopr OID NOT NULL ENCODE RAW' '\n)\nDISTSTYLE EVEN\n;') actual_table_def = self.rs.get_table_definition("pg_catalog.pg_amop") self.assertEqual(expected_table_def, actual_table_def) def test_get_table_definitions(self): expected_table_defs = [{ 'tablename': 'pg_catalog.pg_amop', 'ddl': '--DROP TABLE pg_catalog.pg_amop;' '\nCREATE TABLE IF NOT EXISTS pg_catalog.pg_amop' '\n(' '\n\tamopclaid OID NOT NULL ENCODE RAW' '\n\t,amopsubtype OID NOT NULL ENCODE RAW' '\n\t,amopstrategy SMALLINT NOT NULL ENCODE RAW' '\n\t,amopreqcheck BOOLEAN NOT NULL ENCODE RAW' '\n\t,amopopr OID NOT NULL ENCODE RAW' '\n)\nDISTSTYLE EVEN\n;'}, { 'tablename': 'pg_catalog.pg_amproc', 'ddl': '--DROP TABLE pg_catalog.pg_amproc;' '\nCREATE TABLE IF NOT EXISTS pg_catalog.pg_amproc' '\n(' '\n\tamopclaid OID NOT NULL ENCODE RAW' '\n\t,amprocsubtype OID NOT NULL ENCODE RAW' '\n\t,amprocnum SMALLINT NOT NULL ENCODE RAW' '\n\t,amproc REGPROC NOT NULL ENCODE RAW' '\n)' '\nDISTSTYLE EVEN\n;'}] actual_table_defs = self.rs.get_table_definitions(table="pg_am%p%") self.assertEqual(expected_table_defs, actual_table_defs) def test_get_view_definition(self): expected_view_def = ( '--DROP VIEW pg_catalog.pg_views;' '\nCREATE OR REPLACE VIEW pg_catalog.pg_views AS' '\n SELECT n.nspname AS schemaname' ', c.relname AS viewname' ', pg_get_userbyid(c.relowner) AS viewowner' ', pg_get_viewdef(c.oid) AS definition' '\n FROM pg_class c' '\n LEFT JOIN pg_namespace n ON n.oid = c.relnamespace' '\n WHERE c.relkind = \'v\'::"char";') actual_view_def = self.rs.get_view_definition("pg_catalog.pg_views") self.assertEqual(expected_view_def, actual_view_def) def test_get_view_definitions(self): expected_view_defs = [{ 'viewname': 'pg_catalog.pg_class_info', 'ddl': "--DROP VIEW pg_catalog.pg_class_info;" "\nCREATE OR REPLACE VIEW pg_catalog.pg_class_info AS" "\n SELECT pgc.oid AS reloid, pgc.relname, pgc.relnamespace, " "pgc.reltype, pgc.relowner, pgc.relam, pgc.relfilenode, " "pgc.reltablespace, pgc.relpages, pgc.reltuples, " "pgc.reltoastrelid, pgc.reltoastidxid, pgc.relhasindex, " "pgc.relisshared, pgc.relkind, pgc.relnatts, pgc.relexternid, " "pgc.relisreplicated, pgc.relispinned, pgc.reldiststyle, " "pgc.relprojbaseid, pgc.relchecks, pgc.reltriggers, pgc.relukeys, " "pgc.relfkeys, pgc.relrefs, pgc.relhasoids, pgc.relhaspkey, " "pgc.relhasrules, pgc.relhassubclass, pgc.relacl, " "pgce0.value::smallint AS releffectivediststyle, " "date_add('microsecond'::text, pgce1.value::bigint, " "'2000-01-01 00:00:00'::timestamp without time zone) AS " "relcreationtime" "\n FROM pg_class pgc" "\n LEFT JOIN pg_class_extended pgce0 " "ON pgc.oid = pgce0.reloid AND pgce0.colnum = 0" "\n LEFT JOIN pg_class_extended pgce1 " "ON pgc.oid = pgce1.reloid AND pgce1.colnum = 1;"}] actual_view_def = self.rs.get_view_definitions(view="pg_c%") self.assertEqual(expected_view_defs, actual_view_def) def test_alter_varchar_column_widths(self): append_tbl = Table([['ID', 'Name'], [4, 'Jim'], [5, 'John'], [6, 'Joanna']]) # You can't alter column types if the table has a dependent view self.rs.query(f"DROP VIEW {self.temp_schema}.test_view") # Base table 'Name' column has a width of 5. This should expand it to 6. self.rs.alter_varchar_column_widths(append_tbl, f'{self.temp_schema}.test') self.assertEqual(self.rs.get_columns(self.temp_schema, 'test')['name']['max_length'], 6)