Copy data from Postgres DB to AWS S3 in one step

user_3488109 0
Viewed 4836

I have a use-case to move selected data from Postgres to Amazon S3. This should happen in a single step. I am writing a java program to get this done.

I have figured out a way to copy the data in 2 steps. I use the CopyManager Library and the copyOut method to get the data into my local. After this, I'm moving the same file into S3 using Java.

postgres code to get the data into my local

CopyManager copyManager = new CopyManager((BaseConnection) con);
FileWriter fileWriter = new FileWriter("file.csv");
copyManager.copyOut("COPY (SELECT stmt) TO STDOUT WITH DELIMITER '\t' CSV HEADER", fileWriter);

AWS code to move from local to S3

AmazonS3 conn = new AmazonS3Client(credentials);
conn.setEndpoint("xxx.com");
conn.putObject(
            bucket1.getName(),
            "request.json",
            new File("file.csv")
    );

I expect it to happen in one go instead of writing to a file and then moving the file to S3.

amazon-s3 postgresql java
4 Answers

I've not tried this, but I think you should be able to do it.

Rather than passing a FileWriter to copyOut(), you can pass any OutputStream. You can also provide an InputStream to the putObject method rather than a file.

So you just need to convert your OutputStream to an InputStream, for which there are a number of methods (for example, see this post) or you could use something like the EasyStream library.

user_7018670

You should be using PutObject with InputStream.

You can use PipedOutputStream and PipedInputStream to redirect output to putObject's input

PipedInputStream in = new PipedInputStream();
PipedOutputStream out = new PipedOutputStream(in);

Use in as an argument to PutObject and start writing to out in another thread.

user_53432690

Here is my Python code solution for moving all the tables in a DB named dvdrental in PostgreSQL into an S3 bucket.

import boto3
import json
import psycopg2

s3_client = boto3.client('s3')

# database connection setup
connection = psycopg2.connect(
                database="dvdrental",
                user="postgres",
                password="****")
connection.autocommit = True
cursor = connection.cursor()

# Get all the table names from the database
query = "SELECT table_name FROM information_schema.tables WHERE table_schema = 
'public' and table_type = 'BASE TABLE';"
cursor.execute(query)
table_names = cursor.fetchall()
table_names_list = [row[0] for row in table_names]

# Get the contents of each table in json format and move it to the s3 bucket
for name in table_names_list:
  query2 = "SELECT array_to_json(array_agg(row_to_json(n_alias))) from (select * 
  from {}) n_alias;".format(name)
  cursor.execute(query2)
  file_content = cursor.fetchall()
  data = json.dumps(file_content).encode('UTF-8')
  s3_client.put_object(Body=data, Bucket='dvdrentalbucket', 
  Key='{}.json'.format(name))



# closing database connection.
cursor.close()
connection.close()
print("PostgreSQL connection is closed")
user_166480290

If you're open to doing it in python, here's an example that should work:

import boto
import gzip
import psycopg2
import tempfile

# database connection setup
connection = psycopg2.connect('postgresql://scott:tiger@localhost/mydatabase')
connection.autocommit = True
cursor = connection.cursor()

# aws connection setup
s3_connection = boto.connect_s3('<aws access key>', '<aws secret key>')
bucket = s3_connection.get_bucket('<bucket>')

with tempfile.NamedTemporaryFile() as t:
    with gzip.GzipFile(t.name, mode='wb') as g:
        cursor.copy_expert("COPY ({0}) TO STDOUT WITH CSV HEADER".format('<select_query>'), g)
    key = boto.s3.key.Key(bucket, '<s3_key>')
    key.set_contents_from_filename(g.name)

This process makes use of the tempfile module in python, which allows you to create a file that gets is used and then removed within the process. The context manager (with tempfile...) simplifies the management of the file writing process, so you don't have to delete it manually. Depending on how you set up the tempfile, you can make the file accessible to, or never visible to, system users. Basically, you're streaming the SELECT statement to the STDOUT, and then writing the STDOUT to a tempfile. You're still beholden to your database for the SELECT statement in terms of memory management, speed and access.

The benefit is that you don't need to keep the entire file in memory while trying to transfer it to S3; the drawbacks are that you need enough disk space to temporarily store the file, and that it's obviously slower because you're writing to disk as opposed to doing the entire thing in memory.

The other thing to note is that I kept in the step where python is compressing the file using gzip before uploading. I did that to save space on upload; this is especially useful if you're uploading a table with a lot of repetitive data.

As an aside: you should not use this as-is in an environment where you're open to SQL injection; there are better ways of generating the COPY command if that's a part of your use-case.

user_65918490
Related