1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
|
from io import BytesIO, StringIO
import pytest
@pytest.fixture
def db_table(request, con):
con.run("START TRANSACTION")
con.run(
"CREATE TEMPORARY TABLE t1 "
"(f1 int primary key, f2 int not null, f3 varchar(50) null) "
"on commit drop"
)
return con
def test_copy_to_with_table(db_table):
db_table.run("INSERT INTO t1 (f1, f2, f3) VALUES (:v1, :v1, :v2)", v1=1, v2="1")
db_table.run("INSERT INTO t1 (f1, f2, f3) VALUES (:v1, :v1, :v2)", v1=2, v2="2")
db_table.run("INSERT INTO t1 (f1, f2, f3) VALUES (:v1, :v1, :v2)", v1=3, v2="3")
stream = BytesIO()
db_table.run("copy t1 to stdout", stream=stream)
assert stream.getvalue() == b"1\t1\t1\n2\t2\t2\n3\t3\t3\n"
assert db_table.row_count == 3
def test_copy_to_with_query(con):
stream = BytesIO()
con.run(
"COPY (SELECT 1 as One, 2 as Two) TO STDOUT WITH DELIMITER "
"'X' CSV HEADER QUOTE AS 'Y' FORCE QUOTE Two",
stream=stream,
)
assert stream.getvalue() == b"oneXtwo\n1XY2Y\n"
assert con.row_count == 1
def test_copy_to_with_text_stream(con):
stream = StringIO()
con.run(
"COPY (SELECT 1 as One, 2 as Two) TO STDOUT WITH DELIMITER "
"'X' CSV HEADER QUOTE AS 'Y' FORCE QUOTE Two",
stream=stream,
)
assert stream.getvalue() == "oneXtwo\n1XY2Y\n"
assert con.row_count == 1
def test_copy_from_with_table(db_table):
stream = BytesIO(b"1\t1\t1\n2\t2\t2\n3\t3\t3\n")
db_table.run("copy t1 from STDIN", stream=stream)
assert db_table.row_count == 3
retval = db_table.run("SELECT * FROM t1 ORDER BY f1")
assert retval == [[1, 1, "1"], [2, 2, "2"], [3, 3, "3"]]
def test_copy_from_with_text_stream(db_table):
stream = StringIO("1\t1\t1\n2\t2\t2\n3\t3\t3\n")
db_table.run("copy t1 from STDIN", stream=stream)
retval = db_table.run("SELECT * FROM t1 ORDER BY f1")
assert retval == [[1, 1, "1"], [2, 2, "2"], [3, 3, "3"]]
def test_copy_from_with_query(db_table):
stream = BytesIO(b"f1Xf2\n1XY1Y\n")
db_table.run(
"COPY t1 (f1, f2) FROM STDIN WITH DELIMITER 'X' CSV HEADER "
"QUOTE AS 'Y' FORCE NOT NULL f1",
stream=stream,
)
assert db_table.row_count == 1
retval = db_table.run("SELECT * FROM t1 ORDER BY f1")
assert retval == [[1, 1, None]]
def test_copy_from_with_error(db_table):
stream = BytesIO(b"f1Xf2\n\n1XY1Y\n")
with pytest.raises(BaseException) as e:
db_table.run(
"COPY t1 (f1, f2) FROM STDIN WITH DELIMITER 'X' CSV HEADER "
"QUOTE AS 'Y' FORCE NOT NULL f1",
stream=stream,
)
arg = {
"S": ("ERROR",),
"C": ("22P02",),
"M": (
'invalid input syntax for type integer: ""',
'invalid input syntax for integer: ""',
),
"W": ('COPY t1, line 2, column f1: ""',),
"F": ("numutils.c",),
"R": ("pg_atoi", "pg_strtoint32", "pg_strtoint32_safe"),
}
earg = e.value.args[0]
for k, v in arg.items():
assert earg[k] in v
def test_copy_from_with_text_iterable(db_table):
stream = ["1\t1\t1\n", "2\t2\t2\n", "3\t3\t3\n"]
db_table.run("copy t1 from STDIN", stream=stream)
retval = db_table.run("SELECT * FROM t1 ORDER BY f1")
assert retval == [[1, 1, "1"], [2, 2, "2"], [3, 3, "3"]]
|
