def query_results_to_csv()

in cid/helpers/athena.py [0:0]

• 11 lines of code
• 8 McCabe index (conditional complexity)

    def query_results_to_csv(self, query_id, return_header=False):
        # Get query results
        response = self.client.get_query_results(QueryExecutionId=query_id)

        # Get results rows, either with or without the header row
        rows = response['ResultSet']['Rows'] if return_header else response['ResultSet']['Rows'][1:]

        if rows:
            # Write rows to StringIO in CSV format
            buf = StringIO()
            csv_writer = csv.writer(buf, delimiter=',')

            for row in rows:
                csv_writer.writerow([x['VarCharValue'] if 'VarCharValue' in x else None for x in row['Data']])

            # Strip whitespaces from CSVified string and return it
            return buf.getvalue().rstrip('\n')
        else:
            return None