Skip to content

chore(sinan DAGS): create DAG to fetch dengue data from SINAN #201

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
luabida wants to merge 12 commits into
base: main
Choose a base branch
from
Open

chore(sinan DAGS): create DAG to fetch dengue data from SINAN #201

Changes from 1 commit
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
c0cee37
chore(sinan DAGS): create DAG to fetch dengue data from SINAN
luabida Oct 9, 2023
e940b7a
Include EGH_CONN var to image
luabida Oct 9, 2023
97fb367
fix sql statements
luabida Oct 9, 2023
73f2fec
finish SINAN_DENG DAG
luabida Oct 10, 2023
08e498b
Use parquets chunks, preventing the RAM to get fulfilled
luabida Oct 10, 2023
caebfb6
Handle UndefinedColumn error & add column
luabida Oct 11, 2023
f743f5c
Use recursion to handle to_sql
luabida Oct 11, 2023
8ffb904
Add columns to table before inserting the dataframe
luabida Oct 11, 2023
7d276b5
Parse all columns to TEXT before inserting to db
luabida Oct 11, 2023
d7435c2
minor fixes
luabida Oct 11, 2023
afc99a1
Include SINAN_ZIKA DAG
luabida Oct 13, 2023
c76c431
Include SINAN_CHIK DAG
luabida Oct 13, 2023
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
88 changes: 56 additions & 32 deletions containers/airflow/dags/brasil/sinan/dengue.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ def update_dengue(egh_conn: dict):
import pandas as pd

from sqlalchemy import create_engine, text
from sqlalchemy.exc import ProgrammingError
from pysus.online_data import parquets_to_dataframe
from pysus.ftp.databases.sinan import SINAN

Expand All @@ -49,6 +50,31 @@ def update_dengue(egh_conn: dict):
tablename = "sinan_dengue_m"
files = sinan.get_files(dis_code=dis_code)


def insert_parquets(parquet_dir: str, year: int):
"""
Insert parquet dir into database using its chunks. Delete the chunk
and the directory after insertion.
"""
for parquet in os.listdir(parquet_dir):
file = os.path.join(parquet_dir, parquet)
df = pd.read_parquet(str(file), engine='fastparquet')
df.columns = df.columns.str.lower()
df['year'] = year
df['prelim'] = False
df.to_sql(
name=tablename,
con=create_engine(egh_conn['URI']),
schema="brasil",
if_exists='append',
index=False
)
del df
os.remove(file)
logging.debug(f"{file} inserted into db")
os.rmdir(parquets.path)


f_stage = {}
for file in files:
code, year = sinan.format(file)
Expand Down Expand Up @@ -88,22 +114,21 @@ def update_dengue(egh_conn: dict):

parquets = sinan.download(sinan.get_files(dis_code, year))

for parquet in os.listdir(parquets.path):
file = os.path.join(parquets.path, parquet)
df = pd.read_parquet(str(file), engine='fastparquet')
df.columns = df.columns.str.lower()
df['year'] = year
df['prelim'] = False
df.to_sql(
name=tablename,
con=create_engine(egh_conn['URI']),
schema="brasil",
if_exists='append',
index=False
)
del df
os.remove(file)
logging.debug(f"{file} inserted into db")
try:
insert_parquets(parquets.path, year)
except ProgrammingError as error:
if str(error).startswith("(psycopg2.errors.UndefinedColumn)"):
# Include new columns to table
column_name = str(error).split('"')[1]
Copy link
Copy Markdown
Contributor

@fccoelho fccoelho Oct 11, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think that obtaining the missing column name from the error message is not a good approach, because if psycopg2 changes the wording in their error messages it will break our code. I think we should instead look at the list of column names of the parquet files and compare them with the columns in the current schema. From the difference in these lists, which can be efficiently obtained as list(set(cols1)-set(cols2)), we can then create the alter table query adding the new columns to the database table. With this approach, we don't even need to rely on an exception being raised. This determination of the missing columns can be done before the first insert.

with create_engine(egh_conn['URI']).connect() as conn:
conn.execute(text(
f'ALTER TABLE brasil.{tablename}'
f' ADD COLUMN {column_name} TEXT'
))
conn.commit()
logging.warning(f"Column {column_name} added into {tablename}")
insert_parquets(parquets.path, year)

os.rmdir(parquets.path)

for year in f_stage['prelim']:
Expand All @@ -116,22 +141,21 @@ def update_dengue(egh_conn: dict):

parquets = sinan.download(sinan.get_files(dis_code, year))

for parquet in os.listdir(parquets.path):
file = os.path.join(parquets.path, parquet)
df = pd.read_parquet(str(file), engine='fastparquet')
df.columns = df.columns.str.lower()
df['year'] = year
df['prelim'] = True
df.to_sql(
name=tablename,
con=create_engine(egh_conn['URI']),
schema="brasil",
if_exists='append',
index=False
)
del df
os.remove(file)
logging.debug(f"{file} inserted into db")
try:
insert_parquets(parquets.path, year)
except ProgrammingError as error:
Copy link
Copy Markdown
Contributor

@fccoelho fccoelho Oct 11, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same comment as above

if str(error).startswith("(psycopg2.errors.UndefinedColumn)"):
# Include new columns to table
column_name = str(error).split('"')[1]
with create_engine(egh_conn['URI']).connect() as conn:
conn.execute(text(
f'ALTER TABLE brasil.{tablename}'
f' ADD COLUMN {column_name} TEXT'
))
conn.commit()
logging.warning(f"Column {column_name} added into {tablename}")
insert_parquets(parquets.path, year)

os.rmdir(parquets.path)

update_dengue(CONN)