def test_create_temporary_statement(self):
temp_table = Table(self.tablename, self.columns, self.primary_keys)
expected = ('CREATE TEMPORARY TABLE create_table_test ('
'city VARCHAR(50) NOT NULL, '
'state CHAR(2) NOT NULL, '
'population INTEGER NULL, '
'PRIMARY KEY (city, state));')
result = temp_table.create_temporary_statement()
self.assertSQLStatementEqual(expected, result)
def setUp(self):
self.column_names, self.records = make_records()
self.primary_key_name = self.column_names[0]
self.delete_records = [(self.records[0].city, ),
(self.records[2].city, ),
(self.records[3].city, )]
self.table_name = 'insert_test'
self.columns = [
Column(name='city', data_type='VARCHAR(50)', nullable=False),
Column(name='state', data_type='CHAR(2)', nullable=True)
]
self.primary_key = PrimaryKey([self.primary_key_name])
self.table = Table(self.table_name, self.columns, self.primary_key)
self._setup_table_data()
def setUp(self):
self.schema = 'ddl_schema'
self.tablename = 'create_table_test'
self.qualified_name = 'ddl_schema.create_table_test'
self.column_statement = ("city VARCHAR(50) NOT NULL,"
" state CHAR(2) NOT NULL,"
" population INTEGER NULL,")
self.primary_key_statement = 'PRIMARY KEY (city, state)'
self.columns = table_columns()
self.primary_keys = table_primary_keys()
self.table = Table(self.tablename,
self.columns,
self.primary_keys,
schema=self.schema)
def setUp(self):
self.maxDiff = None
self.schema = 'foo'
self.table_name = 'foobar'
columns = table_columns()
primary_key = table_primary_keys()
self.table = Table(self.table_name, columns, primary_key, self.schema)
class TestDeletePrimaryKeyRecords(PostgresDmlFixture, unittest.TestCase):
def setUp(self):
self.column_names, self.records = make_records()
self.primary_key_name = self.column_names[0]
self.delete_records = [(self.records[0].city, ),
(self.records[2].city, ),
(self.records[3].city, )]
self.table_name = 'insert_test'
self.columns = [
Column(name='city', data_type='VARCHAR(50)', nullable=False),
Column(name='state', data_type='CHAR(2)', nullable=True)
]
self.primary_key = PrimaryKey([self.primary_key_name])
self.table = Table(self.table_name, self.columns, self.primary_key)
self._setup_table_data()
self.delete_processor = dml.DeleteManyPrimaryKey(self.table)
def test_process_delete_insert(self):
self.delete_processor(self.conn, self.delete_records)
expected = set([self.records[1]])
result = set(get_records(self.conn, self.table.qualified_name))
self.assertSetEqual(expected, result)
def _setup_table_data(self):
insert_statement = 'INSERT INTO insert_test (city, state) VALUES (%s, %s)'
with self.conn.cursor() as cursor:
cursor.execute(self.table.create_statement())
for record in self.records:
cursor.execute(insert_statement, record)
self.conn.commit()
def make_temp_copy_table(self):
rand_char = randint(self.RAND_MIN, self.RAND_MAX)
temp_table_name = self._TEMP_FORMATTER.format(
table_name=self.table.name, random=rand_char
)
table_attributes = self.table._asdict()
table_attributes['name'] = temp_table_name
return Table(**table_attributes)
class TestBulkCopy(PostgresDmlFixture, unittest.TestCase):
def setUp(self):
self.table_name = 'insert_record_table'
self.column_names, self.records = make_records()
self.columns = [
Column(self.column_names[0], 'VARCHAR(50)'),
Column(self.column_names[1], 'CHAR(2)', nullable=True)
]
self.primary_key_names = ['city']
self.primary_key = PrimaryKey(self.primary_key_names)
self.table = Table(self.table_name, self.columns, self.primary_key)
self.delimiter = '|'
self.force_null = ['state']
self.null_str = ''
self.insert_query = 'INSERT INTO {} VALUES (%s, %s)'.format(
self.table_name)
with self.conn.cursor() as cursor:
cursor.execute(self.table.create_statement())
self.conn.commit()
@skipPGVersionBefore(*PG_UPSERT_VERSION)
def test_upsert_many(self):
records = [('Miami', 'TX'), ('Chicago', 'MI')]
with self.conn.cursor() as cursor:
cursor.executemany(self.insert_query, records)
self.conn.commit()
bulk_upserter = dml.CopyFromUpsert(self.table,
delimiter=self.delimiter,
null_str=self.null_str,
force_null=self.force_null)
file_object = io.StringIO(delimited_text())
with self.conn:
bulk_upserter(self.conn, file_object)
result = get_records(self.conn, self.table_name)
self.assertEqual(self.records, result)
def test_copy_table_from_csv(self):
self.columns, self.records = make_records()
file_object = io.StringIO(delimited_text())
copy_from_table = dml.CopyFrom(self.table,
delimiter=self.delimiter,
null_str=self.null_str,
force_null=self.force_null)
with self.conn:
copy_from_table(self.conn, file_object)
result = get_records(self.conn, self.table_name)
self.assertEqual(self.records, result)
def reflect_table(conn, table_name, schema='public'):
"""Reflect basic table attributes."""
column_meta = list(get_column_metadata(conn, table_name, schema=schema))
primary_key_columns = list(get_primary_keys(conn, table_name, schema=schema))
columns = [Column(**column_data) for column_data in column_meta]
primary_key = PrimaryKey(primary_key_columns)
return Table(table_name, columns, primary_key, schema=schema)
def setUp(self):
self.table_name = 'insert_record_table'
self.column_names, self.records = make_records()
self.columns = [
Column(self.column_names[0], 'VARCHAR(50)'),
Column(self.column_names[1], 'CHAR(2)', nullable=True)
]
self.primary_key_names = ['city']
self.primary_key = PrimaryKey(self.primary_key_names)
self.table = Table(self.table_name, self.columns, self.primary_key)
self.delimiter = '|'
self.force_null = ['state']
self.null_str = ''
self.insert_query = 'INSERT INTO {} VALUES (%s, %s)'.format(
self.table_name)
with self.conn.cursor() as cursor:
cursor.execute(self.table.create_statement())
self.conn.commit()
class TestBulkCopyAllColumnPrimary(PostgresDmlFixture, unittest.TestCase):
def setUp(self):
self.table_name = 'insert_record_table'
self.column_names, self.records = make_records()
self.records = self.records[0:1]
self.columns = [
Column(self.column_names[0], 'VARCHAR(50)'),
Column(self.column_names[1], 'CHAR(2)')
]
self.primary_key_names = ['city', 'state']
self.primary_key = PrimaryKey(self.primary_key_names)
self.table = Table(self.table_name, self.columns, self.primary_key)
self.delimiter = '|'
self.force_null = []
self.null_str = ''
self.insert_query = 'INSERT INTO {} VALUES (%s, %s)'.format(
self.table_name)
with self.conn.cursor() as cursor:
cursor.execute(self.table.create_statement())
self.conn.commit()
@skipPGVersionBefore(*PG_UPSERT_VERSION)
def test_upsert_many_primary_key(self):
records = [('Chicago', 'IL')]
with self.conn.cursor() as cursor:
cursor.executemany(self.insert_query, records)
self.conn.commit()
bulk_upserter = dml.CopyFromUpsert(self.table,
delimiter=self.delimiter,
null_str=self.null_str,
force_null=self.force_null)
file_object = io.StringIO('\n'.join(
[*delimited_text().splitlines()[0:1], '']))
with self.conn:
bulk_upserter(self.conn, file_object)
result = get_records(self.conn, self.table_name)
self.assertEqual(self.records, result)
@skipPGVersionBefore(*PG_UPSERT_VERSION)
def test_upsert_many_empty_file(self):
bulk_upserter = dml.CopyFromUpsert(self.table,
delimiter=self.delimiter,
null_str=self.null_str,
force_null=self.force_null)
text = '\n'.join([delimited_text().splitlines()[0], ''])
file_object = io.StringIO(text)
with self.conn:
bulk_upserter(self.conn, file_object)
def test_reflect_table(self):
columns = [
Column('mycol', data_type='character(2)', nullable=False),
Column('mycol2', data_type='character(3)', nullable=True)
]
primary_key = PrimaryKey(['mycol'])
expected = Table(self.table, columns, primary_key, schema=self.schema)
result = reflect_table(self.conn, self.table, self.schema)
self.assertEqual(expected, result)
def setUp(self):
self.delete_prefix = 'delete_from__'
self.primary_key_column = Column(name='city',
data_type='VARCHAR(50)',
nullable=False)
self.columns = [
self.primary_key_column,
Column(name='population', data_type='INTEGER', nullable=True)
]
self.primary_key = PrimaryKey(['city'])
self.tablename = 'original_table'
self.schema = 'to_delete'
self.table = Table(self.tablename, self.columns, self.primary_key,
self.schema)
class TestCreateTableEvent(PostgreSQLFixture, unittest.TestCase):
def setUp(self):
self.tablename = 'create_table_event'
self.columns = table_columns()
self.primary_key = table_primary_keys()
self.table_query = "select relname from pg_stat_user_tables where relname=%s;"
self.table = Table(self.tablename, self.columns, self.primary_key)
def test_create_table(self):
with self.conn.cursor() as cursor:
cursor.execute(self.table.create_statement())
self.conn.commit()
with self.conn.cursor() as cursor:
cursor.execute(self.table_query, (self.tablename, ))
table = cursor.fetchone()
self.assertEqual((self.tablename, ), table)
def test_temporary_table(self):
with self.conn as conn:
with conn.cursor() as cursor:
cursor.execute(self.table.create_temporary_statement())
cursor.execute(self.table_query, (self.tablename, ))
temp_table = cursor.fetchone()
cursor.execute(self.table.drop_temporary_statement())
cursor.execute(self.table_query, (self.tablename, ))
no_table = cursor.fetchone()
self.assertEqual((self.tablename, ), temp_table)
self.assertTrue(no_table is None)
def tearDown(self):
with self.conn.cursor() as cursor:
cursor.execute(
'DROP TABLE IF EXISTS {table};'.format(table=self.tablename))
self.conn.commit()
def test_reorder_columns(self):
column_names = ['state', 'city', 'population']
expect_columns = [
Column('state', 'CHAR(2)', False),
Column('city', 'VARCHAR(50)', False),
Column('population', 'INTEGER', True)
]
expect_pkey = PrimaryKey(['state', 'city'])
expected = Table(self.table_name, expect_columns, expect_pkey,
self.schema)
result = order_table_columns(self.table, column_names)
self.assertEqual(expected, result)
class TestDeletePrimaryKeyRecords(PostgresDmlFixture, unittest.TestCase):
def setUp(self):
self.column_names, self.records = make_records()
self.primary_key_name = self.column_names[0]
self.delete_records = [(self.records[0].city, ),
(self.records[2].city, ),
(self.records[3].city, )]
self.table_name = 'insert_test'
self.columns = [
Column(name='city', data_type='VARCHAR(50)', nullable=False),
Column(name='state', data_type='CHAR(2)', nullable=True)
]
self.primary_key = PrimaryKey([self.primary_key_name])
self.table = Table(self.table_name, self.columns, self.primary_key)
self._setup_table_data()
def test_process_delete_insert(self):
delete_processor = dml.DeleteManyPrimaryKey(self.table)
delete_processor(self.conn, self.delete_records)
expected = set([self.records[1]])
result = set(get_records(self.conn, self.table.qualified_name))
self.assertSetEqual(expected, result)
def test_process_delete_copy(self):
text = '\n'.join(
line for index, line in enumerate(delimited_text().split('\n'))
if index != 2)
delete_processor = dml.CopyFromDelete(self.table,
delimiter='|',
header=True)
file_obj = io.StringIO(text)
with self.conn:
delete_processor(self.conn, file_obj)
def _setup_table_data(self):
insert_statement = 'INSERT INTO insert_test (city, state) VALUES (%s, %s)'
with self.conn.cursor() as cursor:
cursor.execute(self.table.create_statement())
for record in self.records:
cursor.execute(insert_statement, record)
self.conn.commit()
def make_temp_copy_table(self):
temp_table_name = self.generate_temp_table_name()
table_attributes = self.table._asdict()
table_attributes['name'] = temp_table_name
return Table(**table_attributes)
class TestTableDDL(PostgresStatementFixture, unittest.TestCase):
def setUp(self):
self.schema = 'ddl_schema'
self.tablename = 'create_table_test'
self.qualified_name = 'ddl_schema.create_table_test'
self.column_statement = ("city VARCHAR(50) NOT NULL,"
" state CHAR(2) NOT NULL,"
" population INTEGER NULL,")
self.primary_key_statement = 'PRIMARY KEY (city, state)'
self.columns = table_columns()
self.primary_keys = table_primary_keys()
self.table = Table(self.tablename,
self.columns,
self.primary_keys,
schema=self.schema)
def test_drop_table_statement(self):
expected = 'DROP TABLE IF EXISTS ddl_schema.create_table_test;'
result = self.table.drop_statement()
self.assertSQLStatementEqual(expected, result)
def test_column_statement(self):
expected = 'city VARCHAR(50) NOT NULL, state CHAR(2) NOT NULL, population INTEGER NULL,'
result = self.table.column_statement
self.assertSQLStatementEqual(expected, result)
def test_primary_key_statement(self):
expected = 'PRIMARY KEY (city, state)'
result = self.table.primary_key_statement
self.assertSQLStatementEqual(expected, result)
def test_primary_key_columns(self):
expected = ['city', 'state']
result = self.table.primary_key_columns
self.assertEqual(expected, result)
def test_column_names(self):
expected = ['city', 'state', 'population']
result = self.table.column_names
self.assertEqual(expected, result)
def test_create_statement(self):
expected = ('CREATE TABLE ddl_schema.create_table_test ('
'city VARCHAR(50) NOT NULL, '
'state CHAR(2) NOT NULL, '
'population INTEGER NULL, '
'PRIMARY KEY (city, state));')
result = self.table.create_statement()
self.assertSQLStatementEqual(expected, result)
def test_create_temporary_statement(self):
temp_table = Table(self.tablename, self.columns, self.primary_keys)
expected = ('CREATE TEMPORARY TABLE create_table_test ('
'city VARCHAR(50) NOT NULL, '
'state CHAR(2) NOT NULL, '
'population INTEGER NULL, '
'PRIMARY KEY (city, state));')
result = temp_table.create_temporary_statement()
self.assertSQLStatementEqual(expected, result)
def test_split_qualified_name(self):
expected = self.schema, self.tablename
result = split_qualified_name(self.qualified_name)
self.assertEqual(expected, result)
expected = 'public', self.tablename
result = split_qualified_name(self.tablename)
self.assertEqual(expected, result)
def setUp(self):
self.tablename = 'create_table_event'
self.columns = table_columns()
self.primary_key = table_primary_keys()
self.table_query = "select relname from pg_stat_user_tables where relname=%s;"
self.table = Table(self.tablename, self.columns, self.primary_key)