Skip to content

Commit 69d0f2e

Browse files
hendrikmakaithendrikmakait
committed
tmpdir
1 parent fc81fe0 commit 69d0f2e

File tree

1 file changed

+89
-42
lines changed

1 file changed

+89
-42
lines changed

‎tests/tpch/test_duckdb.py

+89-42
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,8 @@
1+
import uuid
2+
from collections.abc import Iterator
3+
from contextlib import contextmanager
4+
from pathlib import Path
5+
16
import botocore.session
27
import pytest
38

@@ -8,180 +13,222 @@
813
from . import duckdb_queries # noqa: E402
914

1015

16+
@contextmanager
17+
def tempdir(base_path: str) -> Iterator[str]:
18+
path = Path(base_path) / str(uuid.uuid4())
19+
path.mkdir()
20+
try:
21+
yield path.name
22+
finally:
23+
path.rmdir()
24+
25+
1126
@pytest.fixture
12-
def connection(local, restart, tmp_path):
13-
def _():
14-
con = duckdb.connect()
15-
16-
if not local: # Setup s3 credentials
17-
session = botocore.session.Session()
18-
creds = session.get_credentials()
19-
con.install_extension("httpfs")
20-
con.load_extension("httpfs")
21-
con.sql(
22-
f"""
23-
SET temp_directory='{tmp_path}';
24-
25-
SET s3_region='us-east-2';
26-
SET s3_access_key_id='{creds.access_key}';
27-
SET s3_secret_access_key='{creds.secret_key}';
28-
SET s3_session_token='{creds.token}';
29-
"""
30-
)
31-
return con
27+
def connection(local, restart):
28+
@contextmanager
29+
def _():
30+
if local:
31+
from tempfile import TemporaryDirectory
32+
33+
tempdir_ctx = TemporaryDirectory()
34+
else:
35+
from distributed import get_worker
36+
37+
tempdir_ctx = tempdir(get_worker().local_directory)
38+
39+
with tempdir_ctx as dir, duckdb.connect() as con:
40+
con.sql(f"SET temp_directory='{dir}';")
41+
42+
if not local:
43+
# Setup s3 credentials
44+
session = botocore.session.Session()
45+
creds = session.get_credentials()
46+
con.install_extension("httpfs")
47+
con.load_extension("httpfs")
48+
con.sql(
49+
f"""
50+
SET s3_region='us-east-2';
51+
SET s3_access_key_id='{creds.access_key}';
52+
SET s3_secret_access_key='{creds.secret_key}';
53+
SET s3_session_token='{creds.token}';
54+
"""
55+
)
56+
yield con
3257

3358
return _
3459

3560

3661
def test_query_1(run, connection, dataset_path, scale):
3762
def _():
38-
duckdb_queries.query_1(connection(), dataset_path, scale)
63+
with connection() as con:
64+
duckdb_queries.query_1(con, dataset_path, scale)
3965

4066
run(_)
4167

4268

4369
def test_query_2(run, connection, dataset_path, scale):
4470
def _():
45-
duckdb_queries.query_2(connection(), dataset_path, scale)
71+
with connection() as con:
72+
duckdb_queries.query_2(con, dataset_path, scale)
4673

4774
run(_)
4875

4976

5077
def test_query_3(run, connection, dataset_path, scale):
5178
def _():
52-
duckdb_queries.query_3(connection(), dataset_path, scale)
79+
with connection() as con:
80+
duckdb_queries.query_3(con, dataset_path, scale)
5381

5482
run(_)
5583

5684

5785
def test_query_4(run, connection, dataset_path, scale):
5886
def _():
59-
duckdb_queries.query_4(connection(), dataset_path, scale)
87+
with connection() as con:
88+
duckdb_queries.query_4(con, dataset_path, scale)
6089

6190
run(_)
6291

6392

6493
def test_query_5(run, connection, dataset_path, scale):
6594
def _():
66-
duckdb_queries.query_5(connection(), dataset_path, scale)
95+
with connection() as con:
96+
duckdb_queries.query_5(con, dataset_path, scale)
6797

6898
run(_)
6999

70100

71101
def test_query_6(run, connection, dataset_path, scale):
72102
def _():
73-
duckdb_queries.query_6(connection(), dataset_path, scale)
103+
with connection() as con:
104+
duckdb_queries.query_6(con, dataset_path, scale)
74105

75106
run(_)
76107

77108

78109
def test_query_7(run, connection, dataset_path, scale):
79110
def _():
80-
duckdb_queries.query_7(connection(), dataset_path, scale)
111+
with connection() as con:
112+
duckdb_queries.query_7(con, dataset_path, scale)
81113

82114
run(_)
83115

84116

85117
def test_query_8(run, connection, dataset_path, scale):
86118
def _():
87-
duckdb_queries.query_8(connection(), dataset_path, scale)
119+
with connection() as con:
120+
duckdb_queries.query_8(con, dataset_path, scale)
88121

89122
run(_)
90123

91124

92125
def test_query_9(run, connection, dataset_path, scale):
93126
def _():
94-
duckdb_queries.query_9(connection(), dataset_path, scale)
127+
with connection() as con:
128+
duckdb_queries.query_9(con, dataset_path, scale)
95129

96130
run(_)
97131

98132

99133
def test_query_10(run, connection, dataset_path, scale):
100134
def _():
101-
duckdb_queries.query_10(connection(), dataset_path, scale)
135+
with connection() as con:
136+
duckdb_queries.query_10(con, dataset_path, scale)
102137

103138
run(_)
104139

105140

106141
def test_query_11(run, connection, dataset_path, scale):
107142
def _():
108-
duckdb_queries.query_11(connection(), dataset_path, scale)
143+
with connection() as con:
144+
duckdb_queries.query_11(con, dataset_path, scale)
109145

110146
run(_)
111147

112148

113149
def test_query_12(run, connection, dataset_path, scale):
114150
def _():
115-
duckdb_queries.query_12(connection(), dataset_path, scale)
151+
with connection() as con:
152+
duckdb_queries.query_12(con, dataset_path, scale)
116153

117154
run(_)
118155

119156

120157
def test_query_13(run, connection, dataset_path, scale):
121158
def _():
122-
duckdb_queries.query_13(connection(), dataset_path, scale)
159+
with connection() as con:
160+
duckdb_queries.query_13(con, dataset_path, scale)
123161

124162
run(_)
125163

126164

127165
def test_query_14(run, connection, dataset_path, scale):
128166
def _():
129-
duckdb_queries.query_14(connection(), dataset_path, scale)
167+
with connection() as con:
168+
duckdb_queries.query_14(con, dataset_path, scale)
130169

131170
run(_)
132171

133172

134173
def test_query_15(run, connection, dataset_path, scale):
135174
def _():
136-
duckdb_queries.query_15(connection(), dataset_path, scale)
175+
with connection() as con:
176+
duckdb_queries.query_15(con, dataset_path, scale)
137177

138178
run(_)
139179

140180

141181
def test_query_16(run, connection, dataset_path, scale):
142182
def _():
143-
duckdb_queries.query_16(connection(), dataset_path, scale)
183+
with connection() as con:
184+
duckdb_queries.query_16(con, dataset_path, scale)
144185

145186
run(_)
146187

147188

148189
def test_query_17(run, connection, dataset_path, scale):
149190
def _():
150-
duckdb_queries.query_17(connection(), dataset_path, scale)
191+
with connection() as con:
192+
duckdb_queries.query_17(con, dataset_path, scale)
151193

152194
run(_)
153195

154196

155197
def test_query_18(run, connection, dataset_path, scale):
156198
def _():
157-
duckdb_queries.query_18(connection(), dataset_path, scale)
199+
with connection() as con:
200+
duckdb_queries.query_18(con, dataset_path, scale)
158201

159202
run(_)
160203

161204

162205
def test_query_19(run, connection, dataset_path, scale):
163206
def _():
164-
duckdb_queries.query_19(connection(), dataset_path, scale)
207+
with connection() as con:
208+
duckdb_queries.query_19(con, dataset_path, scale)
165209

166210
run(_)
167211

168212

169213
def test_query_20(run, connection, dataset_path, scale):
170214
def _():
171-
duckdb_queries.query_20(connection(), dataset_path, scale)
215+
with connection() as con:
216+
duckdb_queries.query_20(con, dataset_path, scale)
172217

173218
run(_)
174219

175220

176221
def test_query_21(run, connection, dataset_path, scale):
177222
def _():
178-
duckdb_queries.query_21(connection(), dataset_path, scale)
223+
with connection() as con:
224+
duckdb_queries.query_21(con, dataset_path, scale)
179225

180226
run(_)
181227

182228

183229
def test_query_22(run, connection, dataset_path, scale):
184230
def _():
185-
duckdb_queries.query_22(connection(), dataset_path, scale)
231+
with connection() as con:
232+
duckdb_queries.query_22(con, dataset_path, scale)
186233

187234
run(_)

0 commit comments

Comments
 (0)