Exporting data from PostgreSQL as CSV to S3 bucket I have a general question regarding DRY. As you can see there are several parts such as io.StringIO(), csv.writer, io.BytesIO which are repeated. I now wonder if there is a better way, where I don't seem to repeat myself. Or is the way how I use it here a fair approach and not considered as DRY?

What does your code do?

With the code, I create two CSV-files in the buffer and directly upload them to an AWS S3 bucket. The data for these files are from my Postgres database. The framework I use is Django. That's why I am able to access the database with commands such as Event.objects.filter

What does the rest of your program look like?

That's basically all of the program. After the CSV-files are available on my bucket an AWS Lambda function proceeds with the data and a machine learning model.

What Python version did you write this for?

I use Python 3.

class Forecast:
    @classmethod
    def export_data_for_forecast(cls):
        s3 = boto3.client(
            's3',
            aws_access_key_id=settings.ML_AWS_ACCESS_KEY_ID,
            aws_secret_access_key=settings.ML_AWS_SECRET_ACCESS_KEY,
        )

        BUCKET_NAME = 'fbprophet'
        EVENT_DATA_OBJECT = 'event_data_test.csv'
        ORDER_DATA_OBJECT = 'orders_order_test.csv'

        events = Event.objects.filter(status=EventStatus.LIVE).prefetch_related(
            'orders'
        )

        # Prepare buffer for csv files
        event_data_buffer = io.StringIO()
        event_data_writer = csv.writer(event_data_buffer)
        event_data_writer.writerow(["Event PK", "Name", "Start date"])

        order_data_buffer = io.StringIO()
        order_data_writer = csv.writer(order_data_buffer)
        order_data_writer.writerow(["Event PK", "Created", "Total Gross"])

        active_events = filter(lambda event: not event.is_over, events)
        for event in active_events:
            event_data_writer.writerow([event.pk, event.name, event.start_date])

            for order in event.orders.all():
                order_data_writer.writerow(
                    [order.event.pk, order.created, order.total_gross]
                )

        # Prepare buffer and transform to binary
        event_data_buffer_to_binary = io.BytesIO(
            event_data_buffer.getvalue().encode('utf-8')
        )
        order_data_buffer_to_binary = io.BytesIO(
            order_data_buffer.getvalue().encode('utf-8')
        )

        # Upload to S3
        s3.upload_fileobj(event_data_buffer_to_binary, BUCKET_NAME, EVENT_DATA_OBJECT)
        s3.upload_fileobj(order_data_buffer_to_binary, BUCKET_NAME, ORDER_DATA_OBJECT)
The code's straightforward enough at a glance, I wouldn't have worried
too much about the duplicate bits, especially since it's only double and
not multiple times for each part!  Overdoing abstraction can hurt readability too.  See below for some thoughts on how I'd approach this with that warning in mind.



First of I'd try and move some of the functionality out of the function.
The constants ... should be constants and not right in the middle of the
function body.  Plus, some of them look like constants, while some are
inline ('orders').  I can't tell why the distinction is made here,
maybe there was a reason.

Also consider making all of this not a class method, but a regular one,
then you could for example start creating helper objects beforehand
(like the s3 client) and/or pass in preconstructed ones, e.g. for
testing purposes (like with mock objects).

I'm not going to do these changes right here because you might have had
a reason for the @classmethod.



Instead, let's do a few different changes to make the flow a little bit
easier to comprehend for a reader not already familiar with the code
base:


Moving construction and usage of an object closely together will let
readers just continue reading while not having to jump up the block
again to find out where a particular name was defined ... unless
there are pressing reasons like failing early if no connection could
be made (c.f. the s3 client, though I'm guessing it doesn't actually
establish a connection until the upload_fileobj call is made).
Breaking out functionality into more functions (or even local helper
functions, no one's stoppiing you from creating more abstractions.
Alternatively, and I think that's better suited for this case, create
a new class to wrap things together that belong together.  The hint
for me here is event_data_... appearing with two suffixes and then
order_data_... with the same suffixes.  That begs to be an object
with two attributes.


E.g. like this:

BUCKET_NAME = 'fbprophet'
EVENT_DATA_OBJECT = 'event_data_test.csv'
ORDER_DATA_OBJECT = 'orders_order_test.csv'


class ForecastBuffer:
    def __init__(self, columns):
        # Prepare buffer for csv files
        self.buffer = io.StringIO()
        self.writer = csv.writer(self.buffer)
        self.writer.writerow(columns)

    def write_row(self, row):
        self.writer.writerow(row)

    def upload(self, s3, bucket_name, object_name):
        # Prepare buffer and transform to binary
        to_binary = io.BytesIO(self.buffer.getvalue().encode('utf-8'))
        s3.upload_fileobj(to_binary, bucket_name, object_name)


class Forecast:
    @classmethod
    def export_data_for_forecast(cls):
        event_data = ForecastBuffer(["Event PK", "Name", "Start date"])
        order_data = ForecastBuffer(["Event PK", "Created", "Total Gross"])

        events = Event.objects.filter(status=EventStatus.LIVE).prefetch_related(
            'orders'
        )

        for event in filter(lambda event: not event.is_over, events):
            event_data.write_row([event.pk, event.name, event.start_date])

            for order in event.orders.all():
                order_data.write_row(
                    [order.event.pk, order.created, order.total_gross]
                )

        # Upload to S3
        s3 = boto3.client(
            's3',
            aws_access_key_id=settings.ML_AWS_ACCESS_KEY_ID,
            aws_secret_access_key=settings.ML_AWS_SECRET_ACCESS_KEY,
        )

        event_data.upload(s3, BUCKET_NAME, EVENT_DATA_OBJECT)
        order_data.upload(s3, BUCKET_NAME, ORDER_DATA_OBJECT)


Okay, so, the methods still a bit long in the middle.

One thing that came to my mind was that order.event.pk is probably the
same as event.pk itself, right?  Maybe check and simplify that.

Then, the filter call with the lambda feels a bit weird to me, could
that perhaps already be filtered out in the Event.objects.filter...
chain above it?

You could also move one last bit to make things clearer:  Don't expose
the "raw" writerow call, instead just feed objects to the buffers:

        for event in filter(lambda event: not event.is_over, events):
            event_data.write(event)

            for order in event.orders.all():
                order_data.write(order)


However, that would require to have either a method on the object being
written that specifies the serialisation output (well, which values to
select for the rows), or the same information on the ForecastBuffer
object (you could pass in a function that formats an object for CSV).
Depends entirely on the rest of the code if that's worth it / feasible
at all.



Btw. I just saw that itertools has a complement to filter called
filterfalse (arguably that could be a better name, keep comes to my
mind):

        # from itertools import filterfalse

        for event in filterfalse(Event.is_over, events):
            event_data.write(event)


Assuming that you can get the unbound method Event.is_over that way.



Coming back to that @classmethod, the method here is difficult to test
due to the global variables.  Depending on how it's called it might be
worth to make events a parameter so that the functionality of this
method is simply: format to CSV and upload.  Same goes for settings -
the s3 client most definitely should be passed in so this function
doesn't have to deal with configuration on top of formatting and
uploading.

Actually that's still two things, but it's short enough that that's
probably okay, splitting the formatting into another method might also
be worth it actually:

class Forecast:
    @classmethod
    def format_data_for_forecast(cls, events):
        event_data = ForecastBuffer(["Event PK", "Name", "Start date"])
        order_data = ForecastBuffer(["Event PK", "Created", "Total Gross"])

        for event in filterfalse(Event.is_over, events):
            event_data.write(event)

            for order in event.orders.all():
                order_data.write(order)

        return event_data, order_data

    @classmethod
    def export_data_for_forecast(cls):
        events = Event.objects.filter(status=EventStatus.LIVE).prefetch_related(
            'orders'
        )

        event_data, order_data = cls.format_data_for_forecast(events)

        # Upload to S3
        s3 = boto3.client(
            's3',
            aws_access_key_id=settings.ML_AWS_ACCESS_KEY_ID,
            aws_secret_access_key=settings.ML_AWS_SECRET_ACCESS_KEY,
        )

        event_data.upload(s3, BUCKET_NAME, EVENT_DATA_OBJECT)
        order_data.upload(s3, BUCKET_NAME, ORDER_DATA_OBJECT)