"""This module contains utility functions for connecting to, and updating an in-process DuckDB database instance"""
import duckdb
import logging
from typing import List, Dict
from os import PathLike
logger = logging.getLogger(__name__)
[docs]
def load_target_as_table(
conn: duckdb.DuckDBPyConnection, table_name: str, target_parquet_path: PathLike
) -> None:
"""
Connects to DuckDB using an in-memory database and loads the target Parquet
file into a temporary, mutable table with the given name.
Parameters
----------
conn : duckdb.DuckDBPyConnection
An active DuckDB connection object.
table_name : str
The name of the in-memory table to create.
target_parquet_path : str
The path to the target Parquet file to load.
"""
conn.execute(
f"CREATE OR REPLACE TABLE {table_name} AS SELECT * FROM read_parquet('{target_parquet_path}');"
)
logging.info(f"Loaded '{target_parquet_path}' into in-memory table '{table_name}'.")
[docs]
def get_upsert_counts(
conn: duckdb.DuckDBPyConnection,
target_table_name: str,
update_parquet_path: PathLike,
join_conditions: str,
key_column: str,
) -> Dict[str, int]:
"""
Calculates UPSERT counts (updated and inserted rows) using a single, efficient
LEFT JOIN query on the update file and the registered target table.
Parameters
----------
conn : duckdb.DuckDBPyConnection
An active DuckDB connection object.
target_table_name : str
The name of the temporary DuckDB table containing the target data (t).
update_parquet_path : str
The file path of the update Parquet file (u).
key_column : str
Column name that is used to check matching rows.
Returns
-------
Dict[str, int]
A dictionary containing 'rows_updated', 'rows_inserted', and 'total_incoming'
on success, or None if an exception occurs during calculation.
"""
# Use the first key column from the target table
# to check for a match (IS NOT NULL) or non-match (IS NULL).
key_col_for_check = f"target.{key_column}"
# Combine the match (update) and non-match (insert) counting
# into a single LEFT JOIN query.
combined_count_query = f"""
SELECT
COUNT({key_col_for_check}) AS rows_updated,
COUNT(*) - COUNT({key_col_for_check}) AS rows_inserted,
COUNT(*) AS total_incoming
FROM read_parquet('{update_parquet_path}') AS update
LEFT JOIN {target_table_name} AS target
ON {join_conditions};
"""
logging.info("Executing single LEFT JOIN query for UPSERT counts.")
result = conn.execute(combined_count_query).fetchone()
if result:
counts = {
"rows_updated": result[0],
"rows_inserted": result[1],
"total_incoming": result[2],
}
else:
counts = {"rows_updated": 0, "rows_inserted": 0, "total_incoming": 0}
logging.info(
f"Found {counts['total_incoming']} incoming records in '{update_parquet_path}': "
f"{counts['rows_updated']} will be UPDATED, {counts['rows_inserted']} will be INSERTED."
)
return counts
[docs]
def run_upsert_update(
conn: duckdb.DuckDBPyConnection,
table_name: str,
update_parquet_path: PathLike,
join_conditions: str,
) -> None:
"""
Performs an UPSERT (Update or Insert) operation on a Parquet file
using an in-memory DuckDB instance via the efficient MERGE INTO statement.
The function loads the TARGET file into an in-memory table, merges
the data from the UPDATE file, and then overwrites the original TARGET
file with the final, reconciled table state.
Parameters
----------
conn : duckdb.DuckDBPyConnection
An active DuckDB connection object.
table_name: str
The name of the temporary in-memory table to update (e.g., "intraday_activities").
update_parquet_path: str
The file path to the Parquet file containing the new/update rows.
join_conditions: str
SQL snippet specifying the columns to match on
Returns
-------
None
The function modifies the TARGET Parquet file in place (via overwrite)
and returns no explicit value on successful execution.
"""
# 4. Execute the MERGE INTO statement for UPSERT logic
merge_sql = f"""
MERGE INTO {table_name} AS target
USING read_parquet('{update_parquet_path}') AS update
ON {join_conditions}
WHEN MATCHED THEN
-- Full row replacement: update all columns (*)
UPDATE SET *
WHEN NOT MATCHED THEN
-- Insert new rows: insert all columns (*)
INSERT *;
"""
logging.info(
"Executing MERGE INTO for synchronization operation (Update or Insert)"
)
conn.execute(merge_sql)
logging.info("UPSERT operation complete.")
[docs]
def save_table_to_parquet(
conn: duckdb.DuckDBPyConnection, table_name: str, output_path: PathLike
):
"""Save a specified table to a parquet file
Parameters
----------
conn : duckdb.DuckDBPyConnection
An active DuckDB connection object.
table_name : str
The name of the temporary in-memory table to save
output_path : PathLike
Path to save the table to
"""
export_sql = f"COPY {table_name} TO '{output_path}' (FORMAT PARQUET, OVERWRITE 1);"
conn.execute(export_sql)
logging.info(f"Table: {table_name} saved to '{output_path}' successfully.")
[docs]
def execute_upsert(
table_name: str,
target_parquet_path: PathLike,
update_parquet_path: PathLike,
temp_output_path: PathLike,
key_columns: List[str],
):
"""Overarching function for executing the upsert workflow
Parameters
----------
table_name : str
The name of the temporary in-memory table to load the target parquet file into
target_parquet_path : PathLike
Path to the target parquet file
update_parquet_path : PathLike
Path to the update parquet file
temp_output_path : PathLike
Path to the temp output parquet file
key_columns : List[str]
Key columns to join the tables on
"""
join_conditions = " AND ".join(
[f"target.{col} = update.{col}" for col in key_columns]
)
config = {"enable_external_access": "True"}
with duckdb.connect(database=":memory:", read_only=False, config=config) as conn:
try:
logging.info(
f"Starting Parquet MERGE UPSERT process on target file: '{target_parquet_path} using key columns: {key_columns}'"
)
load_target_as_table(conn, table_name, target_parquet_path)
get_upsert_counts(
conn, table_name, update_parquet_path, join_conditions, key_columns[0]
)
run_upsert_update(conn, table_name, update_parquet_path, join_conditions)
save_table_to_parquet(conn, table_name, temp_output_path)
except duckdb.Error as e:
e.add_note(f"DuckDB error occurred whilst executing upsert: {e}")
raise
except Exception as e:
e.add_note(
f"An unexpected error occurred whilst executing upsert merge: {e}"
)
raise
