build_database.py / Add option to drop EDC tables (#25)

* Drop edc_table when refresh_type='all' or 'reset-table' set to True

* rename 'reset-tables' param to 'drop-tables'

---------

Co-authored-by: jaouenadel <jaouenadel@protonmail.com>

Author: jaouena

pipelines/run.py

@@ -61,7 +61,14 @@ def run():
     type=str,
     help="Comma-separated list of years to process (for custom refresh type)",
 )
-def run_build_database(refresh_type, custom_years):
+@click.option(
+    "--drop-tables",
+    is_flag=True,
+    show_default=True,
+    default=False,
+    help="Drop and re-create edc tables in the database before data insertion.",
+)
+def run_build_database(refresh_type, custom_years, drop_tables):
     """Run build_database task."""
     module = importlib.import_module("tasks.build_database")
     task_func = getattr(module, "execute")
@@ -70,7 +77,11 @@ def run_build_database(refresh_type, custom_years):
     if custom_years:
         custom_years_list = [year.strip() for year in custom_years.split(",")]
 
-    task_func(refresh_type=refresh_type, custom_years=custom_years_list)
+    task_func(
+        refresh_type=refresh_type,
+        custom_years=custom_years_list,
+        drop_tables=drop_tables,
+    )
 
 
 @run.command("download_database")

pipelines/tasks/build_database.py

@@ -114,17 +114,32 @@ def download_extract_insert_yearly_edc_data(year: str):
     return True
 
 
+def drop_edc_tables():
+    """Drop tables using tables names defined in _config_edc.py"""
+    conn = duckdb.connect(DUCKDB_FILE)
+    tables_names = [
+        file_info["table_name"] for file_info in edc_config["files"].values()
+    ]
+    for table_name in tables_names:
+        query = f"DROP TABLE IF EXISTS {table_name};"
+        logger.info(f"Drop table {table_name} (query: {query})")
+        conn.execute(query)
+    return True
+
+
 def process_edc_datasets(
     refresh_type: Literal["all", "last", "custom"] = "last",
     custom_years: List[str] = None,
+    drop_tables: bool = False,
 ):
     """
     Process the EDC datasets.
     :param refresh_type: Refresh type to run
-        - "all": Refresh the data for every possible year
+        - "all": Drop edc tables and import the data for every possible year.
         - "last": Refresh the data only for the last available year
         - "custom": Refresh the data for the years specified in the list custom_years
     :param custom_years: years to update
+    :param drop_tables: Whether to drop edc tables in the database before data insertion.
     :return:
     """
     available_years = edc_config["source"]["available_years"]
@@ -156,6 +171,9 @@ def process_edc_datasets(
 
     logger.info(f"Launching processing of EDC datasets for years: {years_to_update}")
 
+    if drop_tables or (refresh_type == "all"):
+        drop_edc_tables()
+
     for year in years_to_update:
         download_extract_insert_yearly_edc_data(year=year)
 
@@ -164,11 +182,18 @@ def process_edc_datasets(
     return True
 
 
-def execute(refresh_type: str = "all", custom_years: List[str] = None):
+def execute(
+    refresh_type: str = "all",
+    custom_years: List[str] = None,
+    drop_tables: bool = False,
+):
     """
     Execute the EDC dataset processing with specified parameters.
 
     :param refresh_type: Type of refresh to perform ("all", "last", or "custom")
     :param custom_years: List of years to process when refresh_type is "custom"
+    :param drop_tables: Whether to drop edc tables in the database before data insertion.
     """
-    process_edc_datasets(refresh_type=refresh_type, custom_years=custom_years)
+    process_edc_datasets(
+        refresh_type=refresh_type, custom_years=custom_years, drop_tables=drop_tables
+    )