Skip to content

Commit c8b1b04

Browse files
authored
[ez][CH] s32ch script: automatically create a stored data file if not given (#7048)
As in title Reason: I use this script to check that the adapters in the s3 clickhouse replicator lambda work, but it's annoying to keep making stored data files when I usually mean to run the script on everything in the bucket/prefix. In this case, the stored data file doesn't really matter and you want a new one each time This makes it so that a stored data file will automatically get created if one is not provided
1 parent 9c65381 commit c8b1b04

File tree

1 file changed

+13
-2
lines changed

1 file changed

+13
-2
lines changed

tools/rockset_migration/s32ch.py

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import json
1010
import os
1111
import sys
12+
import tempfile
1213
import time
1314
import urllib
1415
from argparse import ArgumentParser
@@ -68,7 +69,6 @@ def parse_args() -> Any:
6869
parser.add_argument(
6970
"--stored-data",
7071
type=str,
71-
required=True,
7272
help=(
7373
"the name of file containing info between runs. ",
7474
"Should be a JSON file with a dict. ",
@@ -138,12 +138,23 @@ def backfill_s3(
138138
bucket: str,
139139
prefix: str,
140140
clickhouse_table: str,
141-
stored_data_file: str,
141+
stored_data_file: Optional[str],
142142
) -> None:
143143
"""
144144
Upload from s3 into clickhouse
145145
"""
146146

147+
# Create a temporary stored data file if none is provided
148+
if stored_data_file is None:
149+
if not os.path.exists("/tmp"):
150+
os.makedirs("/tmp")
151+
temp_fd, stored_data_file = tempfile.mkstemp(
152+
suffix=".json", prefix=f"s3_to_ch__{clickhouse_table}__", dir="/tmp"
153+
)
154+
os.write(temp_fd, "{}\n".encode("utf-8"))
155+
os.close(temp_fd) # Close the file descriptor, we'll use the filename
156+
print(f"Created temporary stored data file: {stored_data_file}")
157+
147158
with open(stored_data_file, "r") as f:
148159
stored_data = json.load(f)
149160

0 commit comments

Comments
 (0)