export/import add tempfile_dir param

Author: yanboer

aws_s3--1.0.0.sql

@@ -49,7 +49,8 @@ CREATE OR REPLACE FUNCTION aws_s3.table_import_from_s3 (
    secret_key text default null,
    session_token text default null,
    endpoint_url text default null,
-   read_timeout integer default 60
+   read_timeout integer default 60,
+   tempfile_dir text default '/var/lib/postgresql/data/'
 ) RETURNS int
 LANGUAGE plpython3u
 AS $$
@@ -112,7 +113,7 @@ AS $$
             body = response['Body']
             user_content_encoding = response.get('x-amz-meta-content-encoding')
 
-            with tempfile.NamedTemporaryFile() as fd:
+            with tempfile.NamedTemporaryFile(dir=tempfile_dir) as fd:
                 if (content_encoding and content_encoding.lower() == 'gzip') or (user_content_encoding and user_content_encoding.lower() == 'gzip'):
                     with gzip.GzipFile(fileobj=body) as gzipfile:
                         while fd.write(gzipfile.read(204800)):
@@ -144,14 +145,15 @@ CREATE OR REPLACE FUNCTION aws_s3.table_import_from_s3(
    s3_info aws_commons._s3_uri_1,
    credentials aws_commons._aws_credentials_1,
    endpoint_url text default null,
-   read_timeout integer default 60
+   read_timeout integer default 60,
+   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) AS num_rows',
-        ['TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'INTEGER']
+        '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']
     )
     return plan.execute(
         [
@@ -182,6 +184,7 @@ CREATE OR REPLACE FUNCTION aws_s3.query_export_to_s3(
     endpoint_url text default null,
     read_timeout integer default 60,
     override boolean default false,
+    tempfile_dir text default '/var/lib/postgresql/data/',
     OUT rows_uploaded bigint,
     OUT files_uploaded bigint,
     OUT bytes_uploaded bigint
@@ -245,7 +248,7 @@ AS $$
             counter += 1
         upload_file_path = get_unique_file_path(base_name, counter, extension)
 
-    with tempfile.NamedTemporaryFile() as fd:
+    with tempfile.NamedTemporaryFile(dir=tempfile_dir) as fd:
         plan = plpy.prepare(
             "COPY ({query}) TO '{filename}' {options}".format(
                 query=query,
@@ -278,15 +281,16 @@ CREATE OR REPLACE FUNCTION aws_s3.query_export_to_s3(
     endpoint_url text default null,
     read_timeout integer default 60,
     override boolean default false,
+    tempfile_dir text default '/var/lib/postgresql/data/',
     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, $11)',
-        ['TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'INTEGER', 'BOOLEAN']
+        'SELECT * FROM aws_s3.query_export_to_s3($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)',
+        ['TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'TEXT', 'INTEGER', 'BOOLEAN', 'TEXT']
     )
     return plan.execute(
         [