def from_postgres(cls, sql, username=None, password=None, host=None, db=None, port=None): """ Args: sql: str A valid SQL statement username: str Required if env variable ``PGUSER`` not populated password: str Required if env variable ``PGPASSWORD`` not populated host: str Required if env variable ``PGHOST`` not populated db: str Required if env variable ``PGDATABASE`` not populated port: int Required if env variable ``PGPORT`` not populated. """ from parsons.databases.postgres import Postgres pg = Postgres(username=username, password=password, host=host, db=db, port=port) return pg.query(sql)
class TestPostgresDB(unittest.TestCase): def setUp(self): self.temp_schema = TEMP_SCHEMA self.pg = Postgres() 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 smallint,name varchar(5)); create view {self.temp_schema}.test_view as (select * from {self.temp_schema}.test); """ # noqa: E501 self.pg.query(setup_sql) self.pg.query(other_sql) def tearDown(self): # Drop the view, the table and the schema teardown_sql = f""" drop schema if exists {self.temp_schema} cascade; """ self.pg.query(teardown_sql) def test_query(self): # Check that query sending back expected result r = self.pg.query('select 1') self.assertEqual(r[0]['?column?'], 1) def test_query_with_parameters(self): table_name = f"{self.temp_schema}.test" self.pg.copy(self.tbl, f"{self.temp_schema}.test", if_exists='append') sql = f"select * from {table_name} where name = %s" name = 'Sarah' r = self.pg.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.pg.query(sql, parameters=names) self.assertEqual(r.num_rows, 2) def test_copy(self): # Copy a table and ensure table exists self.pg.copy(self.tbl, f'{self.temp_schema}.test_copy', if_exists='drop') r = self.pg.query( f"select * from {self.temp_schema}.test_copy where name='Jim'") self.assertEqual(r[0]['id'], 1) # Copy table and ensure truncate works. self.pg.copy(self.tbl, f'{self.temp_schema}.test_copy', if_exists='truncate') tbl = self.pg.query( f"select count(*) from {self.temp_schema}.test_copy") self.assertEqual(tbl.first, 3) # Copy table and ensure that drop works. self.pg.copy(self.tbl, f'{self.temp_schema}.test_copy', if_exists='drop') tbl = self.pg.query( f"select count(*) from {self.temp_schema}.test_copy") self.assertEqual(tbl.first, 3) # Copy table and ensure that append works. self.pg.copy(self.tbl, f'{self.temp_schema}.test_copy', if_exists='append') tbl = self.pg.query( f"select count(*) from {self.temp_schema}.test_copy") self.assertEqual(tbl.first, 6) # Try to copy the table and ensure that default fail works. self.assertRaises(ValueError, self.pg.copy, self.tbl, f'{self.temp_schema}.test_copy') # Try to copy the table and ensure that explicit fail works. self.assertRaises(ValueError, self.pg.copy, self.tbl, f'{self.temp_schema}.test_copy', if_exists='fail') def test_to_postgres(self): self.tbl.to_postgres(f'{self.temp_schema}.test_copy') r = self.pg.query( f"select * from {self.temp_schema}.test_copy where name='Jim'") self.assertEqual(r[0]['id'], 1) def test_from_postgres(self): tbl = Table([['id', 'name'], [1, 'Jim'], [2, 'John'], [3, 'Sarah']]) self.pg.copy(self.tbl, f'{self.temp_schema}.test_copy', if_exists='drop') out_tbl = self.tbl.from_postgres( f"SELECT * FROM {self.temp_schema}.test_copy") assert_matching_tables(out_tbl, tbl)