Ejemplo n.º 1
0
    def to_postgres(self, table_name, username=None, password=None, host=None,
                    db=None, port=None, **copy_args):
        """
        Write a table to a Postgres database.

        Args:
            table_name: str
                The table name and schema (``my_schema.my_table``) to point the file.
            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.
            \**copy_args: kwargs
                See :func:`~parsons.databases.Postgres.copy`` for options.

        Returns:
            ``None``
        """  # noqa: W605

        from parsons.databases.postgres import Postgres
        pg = Postgres(username=username, password=password, host=host, db=db, port=port)
        pg.copy(self, table_name, **copy_args)
Ejemplo n.º 2
0
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)