From 358d3762482047d8b3631750ee289bb58dcddcd9 Mon Sep 17 00:00:00 2001
From: yanboer <yanbo2695388808@gmail.com>
Date: Mon, 27 Mar 2023 14:24:44 +0800
Subject: [PATCH] import add override param

---
 README.md         |  4 ++++
 aws_s3--1.0.0.sql | 17 +++++++++++++----
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 8c710a2..9a91e20 100644
--- a/README.md
+++ b/README.md
@@ -81,6 +81,7 @@ aws_s3.table_import_from_s3 (
    credentials aws_commons._aws_credentials_1,
    endpoint_url text default null,
    read_timeout integer default 60,
+   override boolean default false,
    tempfile_dir text default '/var/lib/postgresql/data/'
 )
 ```
@@ -96,6 +97,7 @@ s3_info | An aws_commons._s3_uri_1 composite type containing the bucket, file pa
 credentials | An aws_commons._aws_credentials_1 composite type containing the access key, secret key, session token credentials. To omit a value, set it `null`.
 endpoint_url | optional endpoint to use (e.g., `http://localhost:4566`)
 read_timeout | The time in seconds till a timeout exception is thrown when attempting to read from a connection.The default is 60 seconds.
+override | Whether to overwrite data. If true, will truncate table.
 tempfile_dir | Specify temporary file location
 
 ##### Example
@@ -175,6 +177,7 @@ aws_s3.table_import_from_s3 (
    session_token text,
    endpoint_url text default null,
    read_timeout integer default 60,
+   override boolean default false,
    tempfile_dir text default '/var/lib/postgresql/data/'
 ) 
 ```
@@ -192,6 +195,7 @@ secret_key | aws secret key
 session_token | optional session token
 endpoint_url | optional endpoint to use (e.g., `http://localhost:4566`)
 read_timeout | The time in seconds till a timeout exception is thrown when attempting to read from a connection.The default is 60 seconds.
+override | Whether to overwrite data. If true, will truncate table.
 tempfile_dir | Specify temporary file location
 
 ##### Example
diff --git a/aws_s3--1.0.0.sql b/aws_s3--1.0.0.sql
index 7720b23..03776ef 100644
--- a/aws_s3--1.0.0.sql
+++ b/aws_s3--1.0.0.sql
@@ -50,6 +50,7 @@ CREATE OR REPLACE FUNCTION aws_s3.table_import_from_s3 (
    session_token text default null,
    endpoint_url text default null,
    read_timeout integer default 60,
+   override boolean default false,
    tempfile_dir text default '/var/lib/postgresql/data/'
 ) RETURNS int
 LANGUAGE plpython3u
@@ -91,6 +92,9 @@ AS $$
         **aws_settings
     )
 
+    if override:
+        plpy.execute("TRUNCATE TABLE {table_name} RESTRICT;".format(table_name=table_name))
+
     formatted_column_list = "({column_list})".format(column_list=column_list) if column_list else ''
     num_rows = 0
 
@@ -146,14 +150,15 @@ CREATE OR REPLACE FUNCTION aws_s3.table_import_from_s3(
    credentials aws_commons._aws_credentials_1,
    endpoint_url text default null,
    read_timeout integer default 60,
+   override boolean default false,
    tempfile_dir text default '/var/lib/postgresql/data/'
 ) RETURNS INT
 LANGUAGE plpython3u
 AS $$
 
     plan = plpy.prepare(
-        'SELECT aws_s3.table_import_from_s3($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) AS num_rows',
-        ['TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'INTEGER', 'TEXT']
+        'SELECT aws_s3.table_import_from_s3($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13) AS num_rows',
+        ['TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'INTEGER', 'BOOLEAN', 'TEXT']
     )
     return plan.execute(
         [
@@ -167,7 +172,9 @@ AS $$
             credentials['secret_key'],
             credentials['session_token'],
             endpoint_url,
-            read_timeout
+            read_timeout,
+            override,
+            tempfile_dir
         ]
     )[0]['num_rows']
 $$;
@@ -306,7 +313,9 @@ AS $$
             credentials.get('session_token') if credentials else None,
             options,
             endpoint_url,
-            read_timeout
+            read_timeout,
+            override,
+            tempfile_dir
         ]
     )
 $$;