From b81cdebc06310b7ca0ba665c72f20229b3c14932 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Far=C3=ADas=20Santana?= Date: Thu, 16 Nov 2023 15:37:17 +0100 Subject: [PATCH] fix: Redshift tests (#18679) * fix: Redshift tests * fix: Update psycopg codecs instead of setting env var * fix: Actually fix tests --- .../test_redshift_batch_export_workflow.py | 8 ++++++-- .../workflows/redshift_batch_export.py | 20 ++++++------------- 2 files changed, 12 insertions(+), 16 deletions(-) diff --git a/posthog/temporal/tests/batch_exports/test_redshift_batch_export_workflow.py b/posthog/temporal/tests/batch_exports/test_redshift_batch_export_workflow.py index 176b487ff94a0..645458719842e 100644 --- a/posthog/temporal/tests/batch_exports/test_redshift_batch_export_workflow.py +++ b/posthog/temporal/tests/batch_exports/test_redshift_batch_export_workflow.py @@ -114,7 +114,7 @@ def redshift_config(): return { "user": user, "password": password, - "database": "dev", + "database": "posthog_batch_exports_test", "schema": "exports_test_schema", "host": host, "port": int(port), @@ -124,7 +124,10 @@ def redshift_config(): @pytest.fixture def postgres_config(redshift_config): """We shadow this name so that setup_postgres_test_db works with Redshift.""" - return redshift_config + psycopg._encodings._py_codecs["UNICODE"] = "utf-8" + psycopg._encodings.py_codecs.update((k.encode(), v) for k, v in psycopg._encodings._py_codecs.items()) + + yield redshift_config @pytest_asyncio.fixture @@ -137,6 +140,7 @@ async def psycopg_connection(redshift_config, setup_postgres_test_db): host=redshift_config["host"], port=redshift_config["port"], ) + connection.prepare_threshold = None yield connection diff --git a/posthog/temporal/workflows/redshift_batch_export.py b/posthog/temporal/workflows/redshift_batch_export.py index 7b008d4f5b674..de342031e9f90 100644 --- a/posthog/temporal/workflows/redshift_batch_export.py +++ b/posthog/temporal/workflows/redshift_batch_export.py @@ -3,7 +3,6 @@ import datetime as dt import itertools import json -import os import typing from dataclasses import dataclass @@ -39,23 +38,16 @@ async def redshift_connection(inputs) -> typing.AsyncIterator[psycopg.AsyncConne This just yields a Postgres connection but we adjust a couple of things required for psycopg to work with Redshift: - 1. Set PGCLIENTENCODING to utf-8 as Redshift reports back UNICODE. + 1. Set UNICODE encoding to utf-8 as Redshift reports back UNICODE. 2. Set prepare_threshold to None on the connection as psycopg attempts to run DEALLOCATE ALL otherwise which is not supported on Redshift. """ - old_value = os.environ.get("PGCLIENTENCODING", None) - os.environ["PGCLIENTENCODING"] = "utf-8" + psycopg._encodings._py_codecs["UNICODE"] = "utf-8" + psycopg._encodings.py_codecs.update((k.encode(), v) for k, v in psycopg._encodings._py_codecs.items()) - try: - async with postgres_connection(inputs) as connection: - connection.prepare_threshold = None - yield connection - - finally: - if old_value is None: - del os.environ["PGCLIENTENCODING"] - else: - os.environ["PGCLIENTENCODING"] = old_value + async with postgres_connection(inputs) as connection: + connection.prepare_threshold = None + yield connection async def insert_records_to_redshift(