A pipeline that moves PostgreSQL data to S3 needs a way to run a query before uploading its results. PostgresHook supplies the PostgreSQL side of that task: initialize it with an Airflow connection ID, then call get_records(query). The method, inherited from DbApiHook, returns a list of tuples. section 8.3.2
Listing 8.12 Example implementation of the PostgresToS3Operator
def execute(self, context):
postgres_hook = PostgresHook(
postgres_conn_id=self._postgres_conn_id
)
s3_hook = S3Hook(aws_conn_id=self._s3_conn_id)
results = postgres_hook.get_records(self._query) #1
s3_hook.load_string( #2
string_data=str(results),
bucket_name=self._s3_bucket,
key=self._s3_key,
)
This minimal operator pairs PostgresHook with S3Hook. It runs a query, then uploads str(results) to an S3 bucket and key. That proves the transfer path, but stringified tuples are not a useful ordinary data-file format. section 8.3.2
Listing 8.13 In-memory conversion of PostgreSQL query results to CSV and upload to S3
def execute(self, context):
postgres_hook = PostgresHook(
postgres_conn_id=self._postgres_conn_id
)
s3_hook = S3Hook(aws_conn_id=self._s3_conn_id)
results = postgres_hook.get_records(self.query)
data_buffer = io.StringIO() #1
csv_writer = csv.writer(data_buffer, lineterminator=os.linesep)
csv_writer.writerows(results)
data_buffer_binary = io.BytesIO(data_buffer.getvalue().encode())
s3_hook.load_file_obj(
file_obj=data_buffer_binary, #2
bucket_name=self._s3_bucket,
key=self._s3_key,
replace=True, #3
)
For a CSV export, write the returned rows with csv.writer into an in-memory StringIO buffer, encode them as BytesIO, and pass the binary buffer to S3Hook.load_file_obj(). This avoids a leftover local file, though the buffer still consumes memory within the Airflow instance. section 8.3.2
PostgresHook can also help verify a database-writing operator. In a test, point its connection lookup at a Docker PostgreSQL container, check that the target table starts empty, run the operator, and query again for inserted data. A pytest-docker-tools fixture manages the container lifecycle; starting it adds a small delay. section 10.2
