#!/usr/bin/python # -*- coding: utf-8 -*- """Test the modern PyGreSQL interface. Sub-tests for the copy methods. Contributed by Christoph Zwerschke. These tests need a database to test against. """ try: import unittest2 as unittest # for Python < 2.7 except ImportError: import unittest try: from collections.abc import Iterable except ImportError: # Python < 3.3 from collections import Iterable import pgdb # the module under test # We need a database to test against. If LOCAL_PyGreSQL.py exists we will # get our information from that. Otherwise we use the defaults. # The current user must have create schema privilege on the database. dbname = 'unittest' dbhost = None dbport = 5432 try: from .LOCAL_PyGreSQL import * except (ImportError, ValueError): try: from LOCAL_PyGreSQL import * except ImportError: pass try: # noinspection PyUnresolvedReferences unicode except NameError: # Python >= 3.0 unicode = str class InputStream: def __init__(self, data): if isinstance(data, unicode): data = data.encode('utf-8') self.data = data or b'' self.sizes = [] def __str__(self): data = self.data if str is unicode: # Python >= 3.0 data = data.decode('utf-8') return data def __len__(self): return len(self.data) def read(self, size=None): if size is None: output, data = self.data, b'' else: output, data = self.data[:size], self.data[size:] self.data = data self.sizes.append(size) return output class OutputStream: def __init__(self): self.data = b'' self.sizes = [] def __str__(self): data = self.data if str is unicode: # Python >= 3.0 data = data.decode('utf-8') return data def __len__(self): return len(self.data) def write(self, data): if isinstance(data, unicode): data = data.encode('utf-8') self.data += data self.sizes.append(len(data)) class TestStreams(unittest.TestCase): def test_input(self): stream = InputStream('Hello, Wörld!') self.assertIsInstance(stream.data, bytes) self.assertEqual(stream.data, b'Hello, W\xc3\xb6rld!') self.assertIsInstance(str(stream), str) self.assertEqual(str(stream), 'Hello, Wörld!') self.assertEqual(len(stream), 14) self.assertEqual(stream.read(3), b'Hel') self.assertEqual(stream.read(2), b'lo') self.assertEqual(stream.read(1), b',') self.assertEqual(stream.read(1), b' ') self.assertEqual(stream.read(), b'W\xc3\xb6rld!') self.assertEqual(stream.read(), b'') self.assertEqual(len(stream), 0) self.assertEqual(stream.sizes, [3, 2, 1, 1, None, None]) def test_output(self): stream = OutputStream() self.assertEqual(len(stream), 0) for chunk in 'Hel', 'lo', ',', ' ', 'Wörld!': stream.write(chunk) self.assertIsInstance(stream.data, bytes) self.assertEqual(stream.data, b'Hello, W\xc3\xb6rld!') self.assertIsInstance(str(stream), str) self.assertEqual(str(stream), 'Hello, Wörld!') self.assertEqual(len(stream), 14) self.assertEqual(stream.sizes, [3, 2, 1, 1, 7]) class TestCopy(unittest.TestCase): cls_set_up = False @staticmethod def connect(): return pgdb.connect(database=dbname, host='%s:%d' % (dbhost or '', dbport or -1)) @classmethod def setUpClass(cls): con = cls.connect() cur = con.cursor() cur.execute("set client_min_messages=warning") cur.execute("drop table if exists copytest cascade") cur.execute("create table copytest (" "id smallint primary key, name varchar(64))") cur.close() con.commit() cur = con.cursor() try: cur.execute("set client_encoding=utf8") cur.execute("select 'Plácido and José'").fetchone() except (pgdb.DataError, pgdb.NotSupportedError): cls.data[1] = (1941, 'Plaacido Domingo') cls.data[2] = (1946, 'Josee Carreras') cls.can_encode = False cur.close() con.close() cls.cls_set_up = True @classmethod def tearDownClass(cls): con = cls.connect() cur = con.cursor() cur.execute("set client_min_messages=warning") cur.execute("drop table if exists copytest cascade") con.commit() con.close() def setUp(self): self.assertTrue(self.cls_set_up) self.con = self.connect() self.cursor = self.con.cursor() self.cursor.execute("set client_encoding=utf8") def tearDown(self): try: self.cursor.close() except Exception: pass try: self.con.rollback() except Exception: pass try: self.con.close() except Exception: pass data = [(1935, 'Luciano Pavarotti'), (1941, 'Plácido Domingo'), (1946, 'José Carreras')] can_encode = True @property def data_text(self): return ''.join('%d\t%s\n' % row for row in self.data) @property def data_csv(self): return ''.join('%d,%s\n' % row for row in self.data) def truncate_table(self): self.cursor.execute("truncate table copytest") @property def table_data(self): self.cursor.execute("select * from copytest") return self.cursor.fetchall() def check_table(self): self.assertEqual(self.table_data, self.data) def check_rowcount(self, number=len(data)): self.assertEqual(self.cursor.rowcount, number) class TestCopyFrom(TestCopy): """Test the copy_from method.""" def tearDown(self): super(TestCopyFrom, self).tearDown() self.setUp() self.truncate_table() super(TestCopyFrom, self).tearDown() def copy_from(self, stream, **options): return self.cursor.copy_from(stream, 'copytest', **options) @property def data_file(self): return InputStream(self.data_text) def test_bad_params(self): call = self.cursor.copy_from call('0\t', 'copytest'), self.cursor call('1\t', 'copytest', format='text', sep='\t', null='', columns=['id', 'name']) self.assertRaises(TypeError, call) self.assertRaises(TypeError, call, None) self.assertRaises(TypeError, call, None, None) self.assertRaises(TypeError, call, '0\t') self.assertRaises(TypeError, call, '0\t', None) self.assertRaises(TypeError, call, '0\t', 42) self.assertRaises(TypeError, call, '0\t', ['copytest']) self.assertRaises(TypeError, call, '0\t', 'copytest', format=42) self.assertRaises(ValueError, call, '0\t', 'copytest', format='bad') self.assertRaises(TypeError, call, '0\t', 'copytest', sep=42) self.assertRaises(ValueError, call, '0\t', 'copytest', sep='bad') self.assertRaises(TypeError, call, '0\t', 'copytest', null=42) self.assertRaises(ValueError, call, '0\t', 'copytest', size='bad') self.assertRaises(TypeError, call, '0\t', 'copytest', columns=42) self.assertRaises(ValueError, call, b'', 'copytest', format='binary', sep=',') def test_input_string(self): ret = self.copy_from('42\tHello, world!') self.assertIs(ret, self.cursor) self.assertEqual(self.table_data, [(42, 'Hello, world!')]) self.check_rowcount(1) def test_input_string_with_newline(self): self.copy_from('42\tHello, world!\n') self.assertEqual(self.table_data, [(42, 'Hello, world!')]) self.check_rowcount(1) def test_input_string_multiple_rows(self): ret = self.copy_from(self.data_text) self.assertIs(ret, self.cursor) self.check_table() self.check_rowcount() if str is unicode: # Python >= 3.0 def test_input_bytes(self): self.copy_from(b'42\tHello, world!') self.assertEqual(self.table_data, [(42, 'Hello, world!')]) self.truncate_table() self.copy_from(self.data_text.encode('utf-8')) self.check_table() else: # Python < 3.0 def test_input_unicode(self): if not self.can_encode: self.skipTest('database does not support utf8') self.copy_from(u'43\tWürstel, Käse!') self.assertEqual(self.table_data, [(43, 'Würstel, Käse!')]) self.truncate_table() self.copy_from(self.data_text.decode('utf-8')) self.check_table() def test_input_iterable(self): self.copy_from(self.data_text.splitlines()) self.check_table() self.check_rowcount() def test_input_iterable_invalid(self): self.assertRaises(IOError, self.copy_from, [None]) def test_input_iterable_with_newlines(self): self.copy_from('%s\n' % row for row in self.data_text.splitlines()) self.check_table() if str is unicode: # Python >= 3.0 def test_input_iterable_bytes(self): self.copy_from(row.encode('utf-8') for row in self.data_text.splitlines()) self.check_table() def test_sep(self): stream = ('%d-%s' % row for row in self.data) self.copy_from(stream, sep='-') self.check_table() def test_null(self): self.copy_from('0\t\\N') self.assertEqual(self.table_data, [(0, None)]) self.assertIsNone(self.table_data[0][1]) self.truncate_table() self.copy_from('1\tNix') self.assertEqual(self.table_data, [(1, 'Nix')]) self.assertIsNotNone(self.table_data[0][1]) self.truncate_table() self.copy_from('2\tNix', null='Nix') self.assertEqual(self.table_data, [(2, None)]) self.assertIsNone(self.table_data[0][1]) self.truncate_table() self.copy_from('3\t') self.assertEqual(self.table_data, [(3, '')]) self.assertIsNotNone(self.table_data[0][1]) self.truncate_table() self.copy_from('4\t', null='') self.assertEqual(self.table_data, [(4, None)]) self.assertIsNone(self.table_data[0][1]) def test_columns(self): self.copy_from('1', columns='id') self.copy_from('2', columns=['id']) self.copy_from('3\tThree') self.copy_from('4\tFour', columns='id, name') self.copy_from('5\tFive', columns=['id', 'name']) self.assertEqual(self.table_data, [ (1, None), (2, None), (3, 'Three'), (4, 'Four'), (5, 'Five')]) self.check_rowcount(5) self.assertRaises(pgdb.ProgrammingError, self.copy_from, '6\t42', columns=['id', 'age']) self.check_rowcount(-1) def test_csv(self): self.copy_from(self.data_csv, format='csv') self.check_table() def test_csv_with_sep(self): stream = ('%d;"%s"\n' % row for row in self.data) self.copy_from(stream, format='csv', sep=';') self.check_table() self.check_rowcount() def test_binary(self): self.assertRaises(IOError, self.copy_from, b'NOPGCOPY\n', format='binary') self.check_rowcount(-1) def test_binary_with_sep(self): self.assertRaises(ValueError, self.copy_from, '', format='binary', sep='\t') def test_binary_with_unicode(self): self.assertRaises(ValueError, self.copy_from, u'', format='binary') def test_query(self): self.assertRaises(ValueError, self.cursor.copy_from, '', "select null") def test_file(self): stream = self.data_file ret = self.copy_from(stream) self.assertIs(ret, self.cursor) self.check_table() self.assertEqual(len(stream), 0) self.assertEqual(stream.sizes, [8192]) self.check_rowcount() def test_size_positive(self): stream = self.data_file size = 7 num_chunks = (len(stream) + size - 1) // size self.copy_from(stream, size=size) self.check_table() self.assertEqual(len(stream), 0) self.assertEqual(stream.sizes, [size] * num_chunks) self.check_rowcount() def test_size_negative(self): stream = self.data_file self.copy_from(stream, size=-1) self.check_table() self.assertEqual(len(stream), 0) self.assertEqual(stream.sizes, [None]) self.check_rowcount() def test_size_invalid(self): self.assertRaises(TypeError, self.copy_from, self.data_file, size='invalid') class TestCopyTo(TestCopy): """Test the copy_to method.""" @classmethod def setUpClass(cls): super(TestCopyTo, cls).setUpClass() con = cls.connect() cur = con.cursor() cur.execute("set client_encoding=utf8") cur.execute("insert into copytest values (%d, %s)", cls.data) cur.close() con.commit() con.close() def copy_to(self, stream=None, **options): return self.cursor.copy_to(stream, 'copytest', **options) @property def data_file(self): return OutputStream() def test_bad_params(self): call = self.cursor.copy_to call(None, 'copytest') call(None, 'copytest', format='text', sep='\t', null='', columns=['id', 'name']) self.assertRaises(TypeError, call) self.assertRaises(TypeError, call, None) self.assertRaises(TypeError, call, None, 42) self.assertRaises(TypeError, call, None, ['copytest']) self.assertRaises(TypeError, call, 'bad', 'copytest') self.assertRaises(TypeError, call, None, 'copytest', format=42) self.assertRaises(ValueError, call, None, 'copytest', format='bad') self.assertRaises(TypeError, call, None, 'copytest', sep=42) self.assertRaises(ValueError, call, None, 'copytest', sep='bad') self.assertRaises(TypeError, call, None, 'copytest', null=42) self.assertRaises(TypeError, call, None, 'copytest', decode='bad') self.assertRaises(TypeError, call, None, 'copytest', columns=42) def test_generator(self): ret = self.copy_to() self.assertIsInstance(ret, Iterable) rows = list(ret) self.assertEqual(len(rows), 3) rows = ''.join(rows) self.assertIsInstance(rows, str) self.assertEqual(rows, self.data_text) self.check_rowcount() if str is unicode: # Python >= 3.0 def test_generator_bytes(self): ret = self.copy_to(decode=False) self.assertIsInstance(ret, Iterable) rows = list(ret) self.assertEqual(len(rows), 3) rows = b''.join(rows) self.assertIsInstance(rows, bytes) self.assertEqual(rows, self.data_text.encode('utf-8')) else: # Python < 3.0 def test_generator_unicode(self): ret = self.copy_to(decode=True) self.assertIsInstance(ret, Iterable) rows = list(ret) self.assertEqual(len(rows), 3) rows = ''.join(rows) self.assertIsInstance(rows, unicode) self.assertEqual(rows, self.data_text.decode('utf-8')) def test_rowcount_increment(self): ret = self.copy_to() self.assertIsInstance(ret, Iterable) for n, row in enumerate(ret): self.check_rowcount(n + 1) def test_decode(self): ret_raw = b''.join(self.copy_to(decode=False)) ret_decoded = ''.join(self.copy_to(decode=True)) self.assertIsInstance(ret_raw, bytes) self.assertIsInstance(ret_decoded, unicode) self.assertEqual(ret_decoded, ret_raw.decode('utf-8')) self.check_rowcount() def test_sep(self): ret = list(self.copy_to(sep='-')) self.assertEqual(ret, ['%d-%s\n' % row for row in self.data]) def test_null(self): data = ['%d\t%s\n' % row for row in self.data] self.cursor.execute('insert into copytest values(4, null)') try: ret = list(self.copy_to()) self.assertEqual(ret, data + ['4\t\\N\n']) ret = list(self.copy_to(null='Nix')) self.assertEqual(ret, data + ['4\tNix\n']) ret = list(self.copy_to(null='')) self.assertEqual(ret, data + ['4\t\n']) finally: self.cursor.execute('delete from copytest where id=4') def test_columns(self): data_id = ''.join('%d\n' % row[0] for row in self.data) data_name = ''.join('%s\n' % row[1] for row in self.data) ret = ''.join(self.copy_to(columns='id')) self.assertEqual(ret, data_id) ret = ''.join(self.copy_to(columns=['id'])) self.assertEqual(ret, data_id) ret = ''.join(self.copy_to(columns='name')) self.assertEqual(ret, data_name) ret = ''.join(self.copy_to(columns=['name'])) self.assertEqual(ret, data_name) ret = ''.join(self.copy_to(columns='id, name')) self.assertEqual(ret, self.data_text) ret = ''.join(self.copy_to(columns=['id', 'name'])) self.assertEqual(ret, self.data_text) self.assertRaises(pgdb.ProgrammingError, self.copy_to, columns=['id', 'age']) def test_csv(self): ret = self.copy_to(format='csv') self.assertIsInstance(ret, Iterable) rows = list(ret) self.assertEqual(len(rows), 3) rows = ''.join(rows) self.assertIsInstance(rows, str) self.assertEqual(rows, self.data_csv) self.check_rowcount(3) def test_csv_with_sep(self): rows = ''.join(self.copy_to(format='csv', sep=';')) self.assertEqual(rows, self.data_csv.replace(',', ';')) def test_binary(self): ret = self.copy_to(format='binary') self.assertIsInstance(ret, Iterable) for row in ret: self.assertTrue(row.startswith(b'PGCOPY\n\377\r\n\0')) break self.check_rowcount(1) def test_binary_with_sep(self): self.assertRaises(ValueError, self.copy_to, format='binary', sep='\t') def test_binary_with_unicode(self): self.assertRaises(ValueError, self.copy_to, format='binary', decode=True) def test_query(self): self.assertRaises(ValueError, self.cursor.copy_to, None, "select name from copytest", columns='noname') ret = self.cursor.copy_to(None, "select name||'!' from copytest where id=1941") self.assertIsInstance(ret, Iterable) rows = list(ret) self.assertEqual(len(rows), 1) self.assertIsInstance(rows[0], str) self.assertEqual(rows[0], '%s!\n' % self.data[1][1]) self.check_rowcount(1) def test_file(self): stream = self.data_file ret = self.copy_to(stream) self.assertIs(ret, self.cursor) self.assertEqual(str(stream), self.data_text) data = self.data_text if str is unicode: # Python >= 3.0 data = data.encode('utf-8') sizes = [len(row) + 1 for row in data.splitlines()] self.assertEqual(stream.sizes, sizes) self.check_rowcount() class TestBinary(TestCopy): """Test the copy_from and copy_to methods with binary data.""" def test_round_trip(self): # fill table from textual data self.cursor.copy_from(self.data_text, 'copytest', format='text') self.check_table() self.check_rowcount() # get data back in binary format ret = self.cursor.copy_to(None, 'copytest', format='binary') self.assertIsInstance(ret, Iterable) data_binary = b''.join(ret) self.assertTrue(data_binary.startswith(b'PGCOPY\n\377\r\n\0')) self.check_rowcount() self.truncate_table() # fill table from binary data self.cursor.copy_from(data_binary, 'copytest', format='binary') self.check_table() self.check_rowcount() if __name__ == '__main__': unittest.main()