Skip to content

Feature/S3 to Postgres optional column order detection #67

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sleblanc23 wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Feature/S3 to Postgres optional column order detection #67

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5a6edc2
add option to detect column names to avoid order issues
sleblanc23 Mar 26, 2025
3dd2f1a
add the same option to the single key function
sleblanc23 Mar 26, 2025
1e5ee9a
add comment to hopefully clarify config
sleblanc23 Mar 26, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions ea_airflow_util/callables/sql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@ def s3_to_postgres(
truncate: bool = False,
delete_qry: Optional[str] = None,
metadata_qry: Optional[str] = None,
column_detection_delimiter: Optional[str] = None, # specify the column delimiter character to use column order detection (i.e. '\t')
**context
):
if column_customization is None:
Expand All @@ -85,6 +86,11 @@ def s3_to_postgres(
print("File at this s3 key is empty.")
raise AirflowSkipException

if column_detection_delimiter:
file_string = s3_hook.read_key(s3_key, s3_bucket)
header = file_string.split('\n')[0].split(column_detection_delimiter)
column_customization = ', '.join(header)

if truncate and not delete_qry:
with conn.cursor() as cur:
logging.info('Truncating table')
Expand Down Expand Up @@ -130,13 +136,19 @@ def s3_dir_to_postgres(
truncate: bool = False,
delete_s3_dir: bool = False,
metadata_qry: Optional[str] = None,
column_detection_delimiter: Optional[str] = None, # specify the column delimiter character to use column order detection (i.e. '\t')
**context
):
s3_hook = S3Hook(s3_conn_id)
s3_creds = s3_hook.get_connection(s3_hook.aws_conn_id)
s3_bucket = s3_creds.schema
s3_keys = s3.list_s3_keys(s3_hook, s3_bucket, s3_key)

if column_detection_delimiter:
file_string = s3_hook.read_key(s3_keys[0], s3_bucket)
header = file_string.split('\n')[0].split(column_detection_delimiter)
column_customization = ', '.join(header)

if truncate:
conn = PostgresHook(pg_conn_id).get_conn()
with conn.cursor() as cur:
Expand Down