def unload_fn_to_test_psycopg2_connect_patch_as_context_manager(config):
with psycopg2.connect(**config) as conn:
with conn.cursor() as cursor:
setup_table_and_insert_data(cursor)
cursor.execute(
UNLOAD_TEMPLATE.format(
COMMAND="UNLOAD",
SELECT_STATEMENT="select * from test_s3_unload_from_redshift",
TO="TO",
LOCATION="s3://mybucket/myfile.csv",
AUTHORIZATION="AUTHORIZATION",
OPTIONAL_ARGS="",
)
)
fetch_values_from_s3_and_assert(cursor, is_gzipped=False)
def copy_fn_to_test_psycopg2_connect_patch_as_context_manager(config):
with psycopg2.connect(**config) as conn:
with conn.cursor() as cursor:
setup_table_and_bucket(cursor)
cursor.execute(
COPY_TEMPLATE.format(
COMMAND="COPY",
LOCATION="s3://mybucket/file.csv",
COLUMNS="",
FROM="from",
CREDENTIALS="credentials",
OPTIONAL_ARGS="",
)
)
fetch_and_assert_psycopg2(cursor)
def copy_fn_to_test_psycopg2_connect_patch(config):
with moto.mock_s3():
conn = psycopg2.connect(**config)
cursor = conn.cursor()
setup_table_and_bucket(cursor)
cursor.execute(
COPY_TEMPLATE.format(
COMMAND="COPY",
LOCATION="s3://mybucket/file.csv",
COLUMNS="(i, f, c, v)",
FROM="from",
CREDENTIALS="credentials",
OPTIONAL_ARGS="EMPTYASNULL",
)
)
fetch_and_assert_psycopg2(cursor)
def mock_psycopg2_connect(*args, **kwargs):
"""Substitute the default cursor with a custom cursor."""
conn = psycopg2.connect(*args, cursor_factory=CustomCursor, **kwargs)
return conn
