diff --git a/posthog/temporal/workflows/postgres_batch_export.py b/posthog/temporal/workflows/postgres_batch_export.py index 1844742880876..3c19c8eae8998 100644 --- a/posthog/temporal/workflows/postgres_batch_export.py +++ b/posthog/temporal/workflows/postgres_batch_export.py @@ -57,7 +57,8 @@ def copy_tsv_to_postgres(tsv_file, postgres_connection, schema: str, table_name: tsv_file.seek(0) with postgres_connection.cursor() as cursor: - cursor.execute(sql.SQL("SET search_path TO {schema}").format(schema=sql.Identifier(schema))) + if schema: + cursor.execute(sql.SQL("SET search_path TO {schema}").format(schema=sql.Identifier(schema))) cursor.copy_from( tsv_file, table_name, @@ -128,6 +129,11 @@ async def insert_into_postgres_activity(inputs: PostgresInsertInputs): ) with postgres_connection(inputs) as connection: with connection.cursor() as cursor: + if inputs.schema: + table_identifier = sql.Identifier(inputs.schema, inputs.table_name) + else: + table_identifier = sql.Identifier(inputs.table_name) + result = cursor.execute( sql.SQL( """ @@ -145,7 +151,7 @@ async def insert_into_postgres_activity(inputs: PostgresInsertInputs): "timestamp" TIMESTAMP WITH TIME ZONE ) """ - ).format(sql.Identifier(inputs.schema, inputs.table_name)) + ).format(table_identifier) ) schema_columns = [