Add a argument to skip interactive question on upload failure

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")

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")

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] "

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] "