From f5f211e46164af4fccd40c86796843bee01c160f Mon Sep 17 00:00:00 2001
From: thedae <thedae@gmail.com>
Date: Thu, 28 Mar 2024 09:24:00 +0100
Subject: [PATCH] Add a argument to skip interactive question on upload failure

---
 raster_loader/cli/bigquery.py  | 8 ++++++++
 raster_loader/cli/snowflake.py | 8 ++++++++
 raster_loader/io/bigquery.py   | 5 +++--
 raster_loader/io/snowflake.py  | 5 +++--
 4 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/raster_loader/cli/bigquery.py b/raster_loader/cli/bigquery.py
index ac9abde..31d8b15 100644
--- a/raster_loader/cli/bigquery.py
+++ b/raster_loader/cli/bigquery.py
@@ -71,6 +71,12 @@ def bigquery(args=None):
     default=False,
     is_flag=True,
 )
+@click.option(
+    "--cleanup-on-failure",
+    help="Clean up resources if the upload fails. Useful for non-interactive scripts.",
+    default=False,
+    is_flag=True,
+)
 @catch_exception()
 def upload(
     file_path,
@@ -84,6 +90,7 @@ def upload(
     chunk_size,
     overwrite=False,
     append=False,
+    cleanup_on_failure=False,
 ):
     from raster_loader.io.common import (
         get_number_of_blocks,
@@ -151,6 +158,7 @@ def upload(
         chunk_size,
         overwrite=overwrite,
         append=append,
+        cleanup_on_failure=cleanup_on_failure,
     )
 
     click.echo("Raster file uploaded to Google BigQuery")
diff --git a/raster_loader/cli/snowflake.py b/raster_loader/cli/snowflake.py
index be73e18..965e531 100644
--- a/raster_loader/cli/snowflake.py
+++ b/raster_loader/cli/snowflake.py
@@ -71,6 +71,12 @@ def snowflake(args=None):
     default=False,
     is_flag=True,
 )
+@click.option(
+    "--cleanup-on-failure",
+    help="Clean up resources if the upload fails. Useful for non-interactive scripts.",
+    default=False,
+    is_flag=True,
+)
 @catch_exception()
 def upload(
     account,
@@ -88,6 +94,7 @@ def upload(
     chunk_size,
     overwrite=False,
     append=False,
+    cleanup_on_failure=False,
 ):
     from raster_loader.io.common import (
         get_number_of_blocks,
@@ -162,6 +169,7 @@ def upload(
         chunk_size,
         overwrite=overwrite,
         append=append,
+        cleanup_on_failure=cleanup_on_failure,
     )
 
     click.echo("Raster file uploaded to Snowflake")
diff --git a/raster_loader/io/bigquery.py b/raster_loader/io/bigquery.py
index 287a392..e995f51 100644
--- a/raster_loader/io/bigquery.py
+++ b/raster_loader/io/bigquery.py
@@ -103,6 +103,7 @@ def upload_raster(
         chunk_size: int = None,
         overwrite: bool = False,
         append: bool = False,
+        cleanup_on_failure: bool = False,
     ):
         """Write a raster file to a BigQuery table."""
         print("Loading raster file to BigQuery...")
@@ -196,7 +197,7 @@ def done_callback(job):
             raise IOError("Error uploading to BigQuery: {}".format(e.message))
 
         except KeyboardInterrupt:
-            delete = ask_yes_no_question(
+            delete = cleanup_on_failure or ask_yes_no_question(
                 "Would you like to delete the partially uploaded table? [yes/no] "
             )
 
@@ -209,7 +210,7 @@ def done_callback(job):
             raise e
 
         except Exception as e:
-            delete = ask_yes_no_question(
+            delete = cleanup_on_failure or ask_yes_no_question(
                 (
                     "Error uploading to BigQuery. "
                     "Would you like to delete the partially uploaded table? [yes/no] "
diff --git a/raster_loader/io/snowflake.py b/raster_loader/io/snowflake.py
index c3ca5f8..9ef5876 100644
--- a/raster_loader/io/snowflake.py
+++ b/raster_loader/io/snowflake.py
@@ -173,6 +173,7 @@ def upload_raster(
         chunk_size: int = None,
         overwrite: bool = False,
         append: bool = False,
+        cleanup_on_failure: bool = False,
     ) -> bool:
         print("Loading raster file to Snowflake...")
 
@@ -240,7 +241,7 @@ def upload_raster(
             raise IOError("Error uploading to Snowflake: {}".format(e.message))
 
         except KeyboardInterrupt:
-            delete = ask_yes_no_question(
+            delete = cleanup_on_failure or ask_yes_no_question(
                 "Would you like to delete the partially uploaded table? [yes/no] "
             )
 
@@ -253,7 +254,7 @@ def upload_raster(
             raise e
 
         except Exception as e:
-            delete = ask_yes_no_question(
+            delete = cleanup_on_failure or ask_yes_no_question(
                 (
                     "Error uploading to Snowflake. "
                     "Would you like to delete the partially uploaded table? [yes/no] "