export_s3 add override param

Author: yanboer

aws_s3--1.0.0.sql

@@ -181,6 +181,7 @@ CREATE OR REPLACE FUNCTION aws_s3.query_export_to_s3(
     options text default null, 
     endpoint_url text default null,
     read_timeout integer default 60,
+    override boolean default false,
     OUT rows_uploaded bigint,
     OUT files_uploaded bigint,
     OUT bytes_uploaded bigint
@@ -233,14 +234,16 @@ AS $$
         **aws_settings
     )
 
-    # generate unique file path
-    file_path_parts = re.match(r'^(.*?)(\.[^.]*$|$)', file_path)
-    base_name = file_path_parts.group(1)
-    extension = file_path_parts.group(2)
-    counter = 0
-    while file_exists(bucket, get_unique_file_path(base_name, counter, extension), s3):
-        counter += 1
-    unique_file_path = get_unique_file_path(base_name, counter, extension)
+    upload_file_path = file_path
+    if not override:
+        # generate unique file path
+        file_path_parts = re.match(r'^(.*?)(\.[^.]*$|$)', upload_file_path)
+        base_name = file_path_parts.group(1)
+        extension = file_path_parts.group(2)
+        counter = 0
+        while file_exists(bucket, get_unique_file_path(base_name, counter, extension), s3):
+            counter += 1
+        upload_file_path = get_unique_file_path(base_name, counter, extension)
 
     with tempfile.NamedTemporaryFile() as fd:
         plan = plpy.prepare(
@@ -261,7 +264,7 @@ AS $$
             num_lines += buffer.count(b'\n')
             size += len(buffer)
         fd.seek(0)
-        s3.upload_fileobj(fd, bucket, unique_file_path)
+        s3.upload_fileobj(fd, bucket, upload_file_path)
         if 'HEADER TRUE' in options.upper():
             num_lines -= 1
         yield (num_lines, 1, size)
@@ -274,15 +277,16 @@ CREATE OR REPLACE FUNCTION aws_s3.query_export_to_s3(
     options text default null, 
     endpoint_url text default null,
     read_timeout integer default 60,
+    override boolean default false,
     OUT rows_uploaded bigint,
     OUT files_uploaded bigint,
     OUT bytes_uploaded bigint
 ) RETURNS SETOF RECORD
 LANGUAGE plpython3u
 AS $$
     plan = plpy.prepare(
-        'SELECT * FROM aws_s3.query_export_to_s3($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)',
-        ['TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'INTEGER']
+        'SELECT * FROM aws_s3.query_export_to_s3($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)',
+        ['TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'INTEGER', 'BOOLEAN']
     )
     return plan.execute(
         [