Skip to content

Commit 9a21a6f

Browse files
dervoetidervoeti
committed
chore: python formatting
1 parent 78966da commit 9a21a6f

2 files changed

Lines changed: 187 additions & 45 deletions

File tree

‎tests/templates/kuttl/smoke/check-s3.py‎

Lines changed: 186 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,11 @@
1010

1111

1212
def get_connection(username, password, namespace):
13-
host = 'trino-coordinator-default-0.trino-coordinator-default.' + namespace + '.svc.cluster.local'
13+
host = (
14+
"trino-coordinator-default-0.trino-coordinator-default."
15+
+ namespace
16+
+ ".svc.cluster.local"
17+
)
1418
# If you want to debug this locally use
1519
# kubectl -n kuttl-test-XXX port-forward svc/trino-coordinator-default 8443
1620
# host = '127.0.0.1'
@@ -19,7 +23,7 @@ def get_connection(username, password, namespace):
1923
host=host,
2024
port=8443,
2125
user=username,
22-
http_scheme='https',
26+
http_scheme="https",
2327
auth=trino.auth.BasicAuthentication(username, password),
2428
session_properties={"query_max_execution_time": "60s"},
2529
)
@@ -34,37 +38,47 @@ def run_query(connection, query):
3438
return cursor.fetchall()
3539

3640

37-
if __name__ == '__main__':
41+
if __name__ == "__main__":
3842
# Construct an argument parser
3943
all_args = argparse.ArgumentParser()
4044
# Add arguments to the parser
41-
all_args.add_argument("-n", "--namespace", required=True, help="Namespace the test is running in")
45+
all_args.add_argument(
46+
"-n", "--namespace", required=True, help="Namespace the test is running in"
47+
)
4248

4349
args = vars(all_args.parse_args())
4450
namespace = args["namespace"]
4551

4652
print("Starting S3 tests...")
4753
connection = get_connection("admin", "admin", namespace)
4854

49-
trino_version = run_query(connection, "select node_version from system.runtime.nodes where coordinator = true and state = 'active'")[0][0]
50-
print(f"[INFO] Testing against Trino version \"{trino_version}\"")
55+
trino_version = run_query(
56+
connection,
57+
"select node_version from system.runtime.nodes where coordinator = true and state = 'active'",
58+
)[0][0]
59+
print(f'[INFO] Testing against Trino version "{trino_version}"')
5160

5261
# Strip SDP release suffix from the version string
53-
trino_product_version = re.split(r'-stackable', trino_version, maxsplit=1)[0]
62+
trino_product_version = re.split(r"-stackable", trino_version, maxsplit=1)[0]
5463

5564
assert len(trino_product_version) >= 3
5665
assert trino_product_version.isnumeric()
5766
assert trino_version == run_query(connection, "select version()")[0][0]
5867

59-
run_query(connection, "CREATE SCHEMA IF NOT EXISTS hive.minio WITH (location = 's3a://trino/')")
68+
run_query(
69+
connection,
70+
"CREATE SCHEMA IF NOT EXISTS hive.minio WITH (location = 's3a://trino/')",
71+
)
6072

6173
run_query(connection, "DROP TABLE IF EXISTS hive.minio.taxi_data")
6274
run_query(connection, "DROP TABLE IF EXISTS hive.minio.taxi_data_copy")
6375
run_query(connection, "DROP TABLE IF EXISTS hive.minio.taxi_data_transformed")
6476
run_query(connection, "DROP TABLE IF EXISTS hive.hdfs.taxi_data_copy")
6577
run_query(connection, "DROP TABLE IF EXISTS iceberg.minio.taxi_data_copy_iceberg")
6678

67-
run_query(connection, """
79+
run_query(
80+
connection,
81+
"""
6882
CREATE TABLE IF NOT EXISTS hive.minio.taxi_data (
6983
vendor_id VARCHAR,
7084
tpep_pickup_datetime VARCHAR,
@@ -77,13 +91,24 @@ def run_query(connection, query):
7791
format = 'csv',
7892
skip_header_line_count = 1
7993
)
80-
""")
81-
assert run_query(connection, "SELECT COUNT(*) FROM hive.minio.taxi_data")[0][0] == 5000
82-
rows_written = run_query(connection, "CREATE TABLE IF NOT EXISTS hive.minio.taxi_data_copy AS SELECT * FROM hive.minio.taxi_data")[0][0]
94+
""",
95+
)
96+
assert (
97+
run_query(connection, "SELECT COUNT(*) FROM hive.minio.taxi_data")[0][0] == 5000
98+
)
99+
rows_written = run_query(
100+
connection,
101+
"CREATE TABLE IF NOT EXISTS hive.minio.taxi_data_copy AS SELECT * FROM hive.minio.taxi_data",
102+
)[0][0]
83103
assert rows_written == 5000 or rows_written == 0
84-
assert run_query(connection, "SELECT COUNT(*) FROM hive.minio.taxi_data_copy")[0][0] == 5000
104+
assert (
105+
run_query(connection, "SELECT COUNT(*) FROM hive.minio.taxi_data_copy")[0][0]
106+
== 5000
107+
)
85108

86-
rows_written = run_query(connection, """
109+
rows_written = run_query(
110+
connection,
111+
"""
87112
CREATE TABLE IF NOT EXISTS hive.minio.taxi_data_transformed AS
88113
SELECT
89114
CAST(vendor_id as BIGINT) as vendor_id,
@@ -93,61 +118,178 @@ def run_query(connection, query):
93118
CAST(trip_distance as DOUBLE) as trip_distance,
94119
CAST(ratecode_id as BIGINT) as ratecode_id
95120
FROM hive.minio.taxi_data
96-
""")[0][0]
121+
""",
122+
)[0][0]
97123
assert rows_written == 5000 or rows_written == 0
98-
assert run_query(connection, "SELECT COUNT(*) FROM hive.minio.taxi_data_transformed")[0][0] == 5000
124+
assert (
125+
run_query(connection, "SELECT COUNT(*) FROM hive.minio.taxi_data_transformed")[
126+
0
127+
][0]
128+
== 5000
129+
)
99130

100131
print("[INFO] Testing HDFS")
101132

102-
run_query(connection, "CREATE SCHEMA IF NOT EXISTS hive.hdfs WITH (location = 'hdfs://hdfs/trino/')")
103-
rows_written = run_query(connection, "CREATE TABLE IF NOT EXISTS hive.hdfs.taxi_data_copy AS SELECT * FROM hive.minio.taxi_data")[0][0]
133+
run_query(
134+
connection,
135+
"CREATE SCHEMA IF NOT EXISTS hive.hdfs WITH (location = 'hdfs://hdfs/trino/')",
136+
)
137+
rows_written = run_query(
138+
connection,
139+
"CREATE TABLE IF NOT EXISTS hive.hdfs.taxi_data_copy AS SELECT * FROM hive.minio.taxi_data",
140+
)[0][0]
104141
assert rows_written == 5000 or rows_written == 0
105-
assert run_query(connection, "SELECT COUNT(*) FROM hive.hdfs.taxi_data_copy")[0][0] == 5000
142+
assert (
143+
run_query(connection, "SELECT COUNT(*) FROM hive.hdfs.taxi_data_copy")[0][0]
144+
== 5000
145+
)
106146

107147
print("[INFO] Testing Iceberg")
108-
run_query(connection, "DROP TABLE IF EXISTS iceberg.minio.taxi_data_copy_iceberg") # Clean up table to don't fail an second run
109-
assert run_query(connection, """
148+
run_query(
149+
connection, "DROP TABLE IF EXISTS iceberg.minio.taxi_data_copy_iceberg"
150+
) # Clean up table to don't fail an second run
151+
assert (
152+
run_query(
153+
connection,
154+
"""
110155
CREATE TABLE IF NOT EXISTS iceberg.minio.taxi_data_copy_iceberg
111156
WITH (partitioning = ARRAY['vendor_id', 'passenger_count'], format = 'parquet')
112157
AS SELECT * FROM hive.minio.taxi_data
113-
""")[0][0] == 5000
158+
""",
159+
)[0][0]
160+
== 5000
161+
)
114162
# Check current count
115-
assert run_query(connection, "SELECT COUNT(*) FROM iceberg.minio.taxi_data_copy_iceberg")[0][0] == 5000
116-
assert run_query(connection, 'SELECT COUNT(*) FROM iceberg.minio."taxi_data_copy_iceberg$snapshots"')[0][0] == 1
117-
assert run_query(connection, 'SELECT COUNT(*) FROM iceberg.minio."taxi_data_copy_iceberg$partitions"')[0][0] == 12
118-
assert run_query(connection, 'SELECT COUNT(*) FROM iceberg.minio."taxi_data_copy_iceberg$files"')[0][0] == 12
163+
assert (
164+
run_query(
165+
connection, "SELECT COUNT(*) FROM iceberg.minio.taxi_data_copy_iceberg"
166+
)[0][0]
167+
== 5000
168+
)
169+
assert (
170+
run_query(
171+
connection,
172+
'SELECT COUNT(*) FROM iceberg.minio."taxi_data_copy_iceberg$snapshots"',
173+
)[0][0]
174+
== 1
175+
)
176+
assert (
177+
run_query(
178+
connection,
179+
'SELECT COUNT(*) FROM iceberg.minio."taxi_data_copy_iceberg$partitions"',
180+
)[0][0]
181+
== 12
182+
)
183+
assert (
184+
run_query(
185+
connection,
186+
'SELECT COUNT(*) FROM iceberg.minio."taxi_data_copy_iceberg$files"',
187+
)[0][0]
188+
== 12
189+
)
119190

120-
assert run_query(connection, "INSERT INTO iceberg.minio.taxi_data_copy_iceberg SELECT * FROM hive.minio.taxi_data")[0][0] == 5000
191+
assert (
192+
run_query(
193+
connection,
194+
"INSERT INTO iceberg.minio.taxi_data_copy_iceberg SELECT * FROM hive.minio.taxi_data",
195+
)[0][0]
196+
== 5000
197+
)
121198

122199
# Check current count
123-
assert run_query(connection, "SELECT COUNT(*) FROM iceberg.minio.taxi_data_copy_iceberg")[0][0] == 10000
124-
assert run_query(connection, 'SELECT COUNT(*) FROM iceberg.minio."taxi_data_copy_iceberg$snapshots"')[0][0] == 2
125-
assert run_query(connection, 'SELECT COUNT(*) FROM iceberg.minio."taxi_data_copy_iceberg$partitions"')[0][0] == 12
126-
assert run_query(connection, 'SELECT COUNT(*) FROM iceberg.minio."taxi_data_copy_iceberg$files"')[0][0] == 24
200+
assert (
201+
run_query(
202+
connection, "SELECT COUNT(*) FROM iceberg.minio.taxi_data_copy_iceberg"
203+
)[0][0]
204+
== 10000
205+
)
206+
assert (
207+
run_query(
208+
connection,
209+
'SELECT COUNT(*) FROM iceberg.minio."taxi_data_copy_iceberg$snapshots"',
210+
)[0][0]
211+
== 2
212+
)
213+
assert (
214+
run_query(
215+
connection,
216+
'SELECT COUNT(*) FROM iceberg.minio."taxi_data_copy_iceberg$partitions"',
217+
)[0][0]
218+
== 12
219+
)
220+
assert (
221+
run_query(
222+
connection,
223+
'SELECT COUNT(*) FROM iceberg.minio."taxi_data_copy_iceberg$files"',
224+
)[0][0]
225+
== 24
226+
)
127227

128-
if trino_version == '377':
228+
if trino_version == "377":
129229
# io.trino.spi.TrinoException: This connector [iceberg] does not support versioned tables
130-
print("[INFO] Skipping the Iceberg tests reading versioned tables for trino version 377 as it does not support versioned tables")
230+
print(
231+
"[INFO] Skipping the Iceberg tests reading versioned tables for trino version 377 as it does not support versioned tables"
232+
)
131233
else:
132234
# Check count for first snapshot
133-
first_snapshot = run_query(connection, 'select snapshot_id from iceberg.minio."taxi_data_copy_iceberg$snapshots" order by committed_at limit 1')[0][0]
134-
assert run_query(connection, f"SELECT COUNT(*) FROM iceberg.minio.taxi_data_copy_iceberg FOR VERSION AS OF {first_snapshot}")[0][0] == 5000
235+
first_snapshot = run_query(
236+
connection,
237+
'select snapshot_id from iceberg.minio."taxi_data_copy_iceberg$snapshots" order by committed_at limit 1',
238+
)[0][0]
239+
assert (
240+
run_query(
241+
connection,
242+
f"SELECT COUNT(*) FROM iceberg.minio.taxi_data_copy_iceberg FOR VERSION AS OF {first_snapshot}",
243+
)[0][0]
244+
== 5000
245+
)
135246

136247
# Compact files
137-
run_query(connection, "ALTER TABLE iceberg.minio.taxi_data_copy_iceberg EXECUTE optimize")
248+
run_query(
249+
connection, "ALTER TABLE iceberg.minio.taxi_data_copy_iceberg EXECUTE optimize"
250+
)
138251

139252
# Check current count
140-
assert run_query(connection, "SELECT COUNT(*) FROM iceberg.minio.taxi_data_copy_iceberg")[0][0] == 10000
141-
assert run_query(connection, 'SELECT COUNT(*) FROM iceberg.minio."taxi_data_copy_iceberg$snapshots"')[0][0] == 3
142-
assert run_query(connection, 'SELECT COUNT(*) FROM iceberg.minio."taxi_data_copy_iceberg$partitions"')[0][0] == 12
143-
assert run_query(connection, 'SELECT COUNT(*) FROM iceberg.minio."taxi_data_copy_iceberg$files"')[0][0] == 12 # Compaction yeah :)
253+
assert (
254+
run_query(
255+
connection, "SELECT COUNT(*) FROM iceberg.minio.taxi_data_copy_iceberg"
256+
)[0][0]
257+
== 10000
258+
)
259+
assert (
260+
run_query(
261+
connection,
262+
'SELECT COUNT(*) FROM iceberg.minio."taxi_data_copy_iceberg$snapshots"',
263+
)[0][0]
264+
== 3
265+
)
266+
assert (
267+
run_query(
268+
connection,
269+
'SELECT COUNT(*) FROM iceberg.minio."taxi_data_copy_iceberg$partitions"',
270+
)[0][0]
271+
== 12
272+
)
273+
assert (
274+
run_query(
275+
connection,
276+
'SELECT COUNT(*) FROM iceberg.minio."taxi_data_copy_iceberg$files"',
277+
)[0][0]
278+
== 12
279+
) # Compaction yeah :)
144280

145281
# Test could be improved by also testing update and deletes
146282

147283
# Test postgres connection
148-
run_query(connection, 'SHOW SCHEMAS IN postgresgeneric')
149-
run_query(connection, 'CREATE SCHEMA IF NOT EXISTS postgresgeneric.tpch')
150-
run_query(connection, 'CREATE TABLE IF NOT EXISTS postgresgeneric.tpch.nation AS SELECT * FROM tpch.tiny.nation')
151-
assert run_query(connection, "SELECT COUNT(*) FROM postgresgeneric.tpch.nation")[0][0] == 25
284+
run_query(connection, "SHOW SCHEMAS IN postgresgeneric")
285+
run_query(connection, "CREATE SCHEMA IF NOT EXISTS postgresgeneric.tpch")
286+
run_query(
287+
connection,
288+
"CREATE TABLE IF NOT EXISTS postgresgeneric.tpch.nation AS SELECT * FROM tpch.tiny.nation",
289+
)
290+
assert (
291+
run_query(connection, "SELECT COUNT(*) FROM postgresgeneric.tpch.nation")[0][0]
292+
== 25
293+
)
152294

153295
print("[SUCCESS] All tests in check-s3.py succeeded!")

‎tests/templates/kuttl/smoke_aws/check-s3.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ def run_query(connection, query):
5959
print(f'[INFO] Testing against Trino version "{trino_version}"')
6060

6161
# Strip SDP release suffix from the version string
62-
trino_product_version = re.split(r'-stackable', trino_version, maxsplit=1)[0]
62+
trino_product_version = re.split(r"-stackable", trino_version, maxsplit=1)[0]
6363

6464
assert len(trino_product_version) >= 3
6565
assert trino_product_version.isnumeric()

0 commit comments

Comments
 (0)