def test_delimiter(): tablename = "delimiter" asset = "tests/assets/delimiter.csv" copy_to( HOST, PORT, DBNAME, USER, PASSWORD, tablename, asset, overwrite=True, delimiter="@", ) with psycopg2.connect(DSN) as conn: with conn.cursor() as curs: curs.execute( "SELECT * FROM {tablename}".format(tablename=tablename)) rows = curs.fetchall() assert len(rows) == 10 for row in rows: assert len(row) == 3 for col in row: assert col is not None
def test_check_delimiter(): tablename = "with_delimiter_error" asset = "tests/assets/error_delimiter.csv" copy_to( HOST, PORT, DBNAME, USER, PASSWORD, tablename, asset, overwrite=True, skip_error=True, ) with open(asset + ".err") as f: errors = f.readlines() assert len(errors) == 3 assert errors[1].startswith("2") assert errors[1].split(",")[1].startswith("MissingFieldsException") assert errors[2].startswith("5") assert errors[2].split(",")[1].startswith("TooManyFieldsException") with psycopg2.connect(DSN) as conn: with conn.cursor() as curs: curs.execute( "SELECT * FROM {tablename}".format(tablename=tablename)) rows = curs.fetchall() assert len(rows) == 8 assert 2 not in [a for a, b, c in rows] assert 5 not in [a for a, b, c in rows]
def test_check_unterminated_quote(): tablename = "with_quote_error" asset = "tests/assets/error_unterminated_quote.csv" copy_to( HOST, PORT, DBNAME, USER, PASSWORD, tablename, asset, overwrite=True, skip_error=True, ) with open(asset + ".err") as f: errors = f.readlines() assert len(errors) == 2 assert errors[1].startswith("3") assert errors[1].split(",")[1].startswith("WrongFieldDialectException") with psycopg2.connect(DSN) as conn: with conn.cursor() as curs: curs.execute( "SELECT * FROM {tablename}".format(tablename=tablename)) rows = curs.fetchall() assert len(rows) == 9 assert 3 not in [a for a, b, c in rows]
def test_rownum_no_header(): tablename = "rownum_no_header" asset = "tests/assets/no_header.csv" copy_to( HOST, PORT, DBNAME, USER, PASSWORD, tablename, asset, overwrite=True, delimiter="@", header=False, inject_rownum=True, ) with psycopg2.connect(DSN) as conn: with conn.cursor() as curs: curs.execute( "SELECT * FROM {tablename}".format(tablename=tablename)) rows = curs.fetchall() assert len(rows) == 10 for i, row in enumerate(rows): assert len(row) == 4 assert int(row[0]) == (i + 1) for col in row: assert col is not None
def test_rownum_encoding_GB18030(): tablename = "rownum_encoding_GB18030" asset = "tests/assets/encoding_GB18030.csv" copy_to( HOST, PORT, DBNAME, USER, PASSWORD, tablename, asset, overwrite=True, inject_rownum=True, encoding="GB18030", ) with psycopg2.connect(DSN) as conn: with conn.cursor() as curs: curs.execute( "SELECT * FROM {tablename}".format(tablename=tablename)) rows = curs.fetchall() assert len(rows) == 1 assert rows[0][0] == 1 assert rows[0][1] == "气广"
def test_nulls_custom(): tablename = "nulls_custom" asset = "tests/assets/nulls_custom.csv" copy_to( HOST, PORT, DBNAME, USER, PASSWORD, tablename, asset, overwrite=True, delimiter=":", null="NULLZZ", ) with psycopg2.connect(DSN) as conn: with conn.cursor() as curs: curs.execute( "SELECT * FROM {tablename}".format(tablename=tablename)) rows = curs.fetchall() assert len(rows) == 10 for row in rows: assert len(row) == 3 for i, col in enumerate(row): if i == 1: assert col == "" else: assert col is None
def test_error_unterminated_quote(): tablename = "error_unterminated_quote" asset = "tests/assets/error_unterminated_quote.csv" with pytest.raises(psycopg2.errors.BadCopyFileFormat): copy_to(HOST, PORT, DBNAME, USER, PASSWORD, tablename, asset) with psycopg2.connect(DSN) as conn: with conn.cursor() as curs: curs.execute( "SELECT * FROM {tablename}".format(tablename=tablename)) rows = curs.fetchall() assert len(rows) == 0
def test_append(): tablename = "append" asset = "tests/assets/simple.csv" copy_to(HOST, PORT, DBNAME, USER, PASSWORD, tablename, asset) copy_to(HOST, PORT, DBNAME, USER, PASSWORD, tablename, asset) with psycopg2.connect(DSN) as conn: with conn.cursor() as curs: curs.execute( "SELECT * FROM {tablename}".format(tablename=tablename)) rows = curs.fetchall() assert len(rows) == 20
def test_insert(): tablename = "insert" asset = "tests/assets/simple.csv" copy_to(HOST, PORT, DBNAME, USER, PASSWORD, tablename, asset) with psycopg2.connect(DSN) as conn: with conn.cursor() as curs: curs.execute( "SELECT * FROM {tablename}".format(tablename=tablename)) rows = curs.fetchall() assert len(rows) == 10 for row in rows: assert len(row) == 3 for col in row: assert col is not None
def test_unlogged(): tablename = "not_logged" asset = "tests/assets/simple.csv" copy_to( HOST, PORT, DBNAME, USER, PASSWORD, tablename, asset, overwrite=True, unlogged=True, ) with psycopg2.connect(DSN) as conn: with conn.cursor() as curs: curs.execute( "SELECT relowner FROM pg_class WHERE relpersistence = 'u' AND relname = '{tablename}'" .format(tablename=tablename)) rows = curs.fetchone() assert rows[0] == 10
def test_encoding_iso_8859_1(): tablename = "encoding_iso_8859_1" asset = "tests/assets/encoding_ISO_8859_1.csv" copy_to( HOST, PORT, DBNAME, USER, PASSWORD, tablename, asset, overwrite=True, encoding="ISO-8859-1", ) with psycopg2.connect(DSN) as conn: with conn.cursor() as curs: curs.execute( "SELECT * FROM {tablename}".format(tablename=tablename)) rows = curs.fetchall() assert len(rows) == 1 assert rows[0][0] == "ÑÆÛý¼àáäæ"
def test_filename(): tablename = "with_filename" asset = "tests/assets/simple.csv" copy_to( HOST, PORT, DBNAME, USER, PASSWORD, tablename, asset, overwrite=True, inject_filename=True, ) with psycopg2.connect(DSN) as conn: with conn.cursor() as curs: curs.execute( "SELECT * FROM {tablename}".format(tablename=tablename)) rows = curs.fetchall() assert len(rows) == 10 for c in rows: assert c[0] == "simple.csv"
def test_complex(): tablename = "complex" asset = "tests/assets/complex_LE-LF_header_bom.csv" copy_to( HOST, PORT, DBNAME, USER, PASSWORD, tablename, asset, overwrite=True, lineterminator="\n", ) with psycopg2.connect(DSN) as conn: with conn.cursor() as curs: curs.execute( "SELECT * FROM {tablename}".format(tablename=tablename)) rows = curs.fetchall() assert len(rows) == 1000 for row in rows: assert len(row) == 19
def cli( hostname, port, dbname, username, password, verbose, progress, skip_error, header, rownum, filename, delimiter, quotechar, doublequote, escapechar, lineterminator, null, encoding, overwrite, unlogged, buffer, table, filepath, ): """ COPY FROM 'csv' TO 'postgres' """ default_options = { "application_name": "csv2pg", "connect_timeout": "10", } pgpassword = os.getenv("PGPASSWORD") if password: pgpassword = click.prompt("Password", hide_input=True) copy_to( hostname, port, dbname, username, pgpassword, table, filepath, connection_options=default_options, verbose=verbose, progress=progress, skip_error=skip_error, header=header, inject_rownum=rownum, inject_filename=filename, delimiter=delimiter, quotechar=quotechar, doublequote=doublequote, escapechar=escapechar, lineterminator=lineterminator, null=null, encoding=encoding, overwrite=overwrite, unlogged=unlogged, buffer=buffer, )