Skip to content

Commit 941548a

Browse files
committed
chg(pit_registering_salmon): add last succesfull run + fixing some linter issues
1 parent 61eedff commit 941548a

1 file changed

Lines changed: 63 additions & 4 deletions

File tree

src/datasync/pit_registering_salmon.py

Lines changed: 63 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,56 @@
4646
}
4747

4848

49+
def write_timestamp(
50+
con: duckdb.DuckDBPyConnection,
51+
bucket: str,
52+
prefix: str,
53+
) -> str:
54+
"""
55+
TODO: This is also used in NVA script, move to a common utils
56+
Write last successful run timestamp to S3.
57+
"""
58+
timestamp = datetime.now().isoformat()
59+
con.execute(f"""
60+
COPY (SELECT '{timestamp}' as last_successful_run)
61+
TO 's3://{bucket}/{prefix}/last_successful_run.parquet'
62+
(FORMAT PARQUET, COMPRESSION ZSTD)
63+
""")
64+
log.info("Last successful run timestamp written", time=timestamp)
65+
return timestamp
66+
67+
68+
def setup_duckdb_s3_connection(
69+
endpoint_url: str,
70+
access_key: str,
71+
secret_key: str,
72+
region: str,
73+
) -> duckdb.DuckDBPyConnection:
74+
"""
75+
TODO: This is also used in NVA script, move to a common utils
76+
Set up DuckDB connection with S3 support
77+
"""
78+
con = duckdb.connect()
79+
con.execute("INSTALL httpfs;")
80+
con.execute("LOAD httpfs;")
81+
82+
# Clean endpoint URL
83+
clean_endpoint = endpoint_url.replace("https://", "").replace("http://", "")
84+
85+
con.execute(f"""
86+
CREATE OR REPLACE SECRET (
87+
TYPE S3,
88+
KEY_ID '{access_key}',
89+
SECRET '{secret_key}',
90+
ENDPOINT '{clean_endpoint}',
91+
REGION '{region}',
92+
URL_STYLE 'path'
93+
);
94+
""")
95+
96+
return con
97+
98+
4999
def hex_to_decimal_tag(hex_tag):
50100
"""
51101
Convert hexadecimal PIT tag format to ISO decimal format.
@@ -105,7 +155,7 @@ def get_environmental_data(
105155
try:
106156
yield from client.paginate(
107157
f"enviro/{location_code}",
108-
method="get",
158+
method="GET",
109159
params={
110160
"begin_dt": begin_date,
111161
"end_dt": end_date,
@@ -148,7 +198,7 @@ def get_tags_data(
148198
try:
149199
yield from client.paginate(
150200
f"tags/{location_code}",
151-
method="get",
201+
method="GET",
152202
params={
153203
"begin_dt": begin_date,
154204
"end_dt": end_date,
@@ -192,7 +242,7 @@ def get_readers_voltage_data(
192242
try:
193243
yield from client.paginate(
194244
f"reader/{location_code}",
195-
method="get",
245+
method="GET",
196246
params={
197247
"begin_dt": begin_date,
198248
"end_dt": end_date,
@@ -481,7 +531,7 @@ def replicate(
481531
if tags and check_table_has_data(duckdb_path, "tags", dataset_name):
482532
stream_configs["tags"] = {
483533
"table_name": f"{dataset_name}.tags",
484-
"primary_key": ["detected_at"],
534+
"primary_key": ["detected_at", "tag"],
485535
"update_key": "detected_at",
486536
"time_column": "detected_at",
487537
"location": "antenna__reader__site__slug",
@@ -532,6 +582,15 @@ def replicate(
532582
debug=True,
533583
).run()
534584

585+
con = setup_duckdb_s3_connection(
586+
endpoint_url=endpoint_url,
587+
access_key=access_key,
588+
secret_key=secret_key,
589+
region=region,
590+
)
591+
592+
write_timestamp(con, bucket, "tables")
593+
535594

536595
if __name__ == "__main__":
537596
app()

0 commit comments

Comments
 (0)