chore(DAG-SINAN): proposal to split DAGs by disease (#211)

* chore(DAG-SINAN): proposal to split DAGs by disease

* Changing nomenclature to follow EGH table pattern

* change schemaname back to 'brasil'

* Fix \x00 error on SIFG

* Linter

* Fix docstrings

* Metadata sheets

* Fixing metadata files

* extract.medatada method

* Unittest for metadata

* Extracting all metadata & types from xlsx files

* Linter

* Poetry lock

* Passing files to PySUS

* Move changes to PySUS

* Minor fixes

* SQLAlchemy 2.0 breaks airflow

* Poetry lock

Author: luabida

conda/dev.yaml

@@ -8,7 +8,7 @@ dependencies:
   - nodejs
   - pandas
   - pip
-  - poetry
+  - poetry >= 1.3.2
   - psycopg2
   - python 3.9.*
   - shellcheck

epigraphhub/data/brasil/sinan/__init__.py

@@ -0,0 +1,16 @@
+import unicodedata
+
+from pysus.online_data import SINAN
+
+DISEASES = SINAN.agravos
+
+
+def normalize_str(disease: str) -> str:
+    """
+    Animais Peçonhentos -> animais_peconhentos
+    """
+    non_ascii = (
+        unicodedata.normalize("NFKD", disease).encode("ascii", "ignore").decode()
+    )
+    disease = non_ascii.lower().replace(" ", "_")
+    return disease

epigraphhub/data/brasil/sinan/extract.py

@@ -1,29 +1,34 @@
-from pathlib import PosixPath
+import os
+from pathlib import Path
 
+import pandas as pd
 from loguru import logger
-from pysus.online_data import SINAN
+from pysus import SINAN
 
 from epigraphhub.data._config import PYSUS_DATA_PATH, SINAN_LOG_PATH
 
 logger.add(SINAN_LOG_PATH, retention="7 days")
 
-diseases = SINAN.agravos
 
-
-def download(disease: str):
+def download(disease: str, years: list = None) -> None:
     """
-    Download all parquets available for an disease,
+    Download all parquets available for a disease,
     according to `SINAN.agravos`.
 
     Attrs:
         disease (str): The disease to be downloaded.
-        data_dir (str) : The output directory were files will be downloaded.
-                         A directory with the disease code will be created.
-
-    Returns:
-        parquets_paths_list list(PosixPath) : A list with all parquets dirs.
     """
 
-    SINAN.download_all_years_in_chunks(disease, data_dir=PYSUS_DATA_PATH)
+    SINAN.download_parquets(disease, years, data_path=PYSUS_DATA_PATH)
 
     logger.info(f"All years for {disease} downloaded at {PYSUS_DATA_PATH}")
+
+
+def metadata_df(disease: str) -> pd.DataFrame:
+    """
+    Returns a DataFrame containing metadata for a SINAN disease.
+    """
+    try:
+        return SINAN.metadata_df(disease)
+    except Exception:
+        logger.error(f"Metadata not available for {disease}")

epigraphhub/data/brasil/sinan/loading.py

@@ -1,55 +1,62 @@
 import os
 from pathlib import Path
 
+import pandas as pd
+from pysus import SINAN
 from loguru import logger
 from pangres import upsert
-from pysus.online_data import parquets_to_dataframe as to_df
+from pysus.classes.sinan import Disease
 
 from epigraphhub.connection import get_engine
-from epigraphhub.data._config import PYSUS_DATA_PATH, SINAN_LOG_PATH
+from epigraphhub.data._config import SINAN_LOG_PATH, PYSUS_DATA_PATH
 from epigraphhub.settings import env
 
-logger.add(SINAN_LOG_PATH, retention="7 days")
+from . import normalize_str
 
+logger.add(SINAN_LOG_PATH, retention="7 days")
 engine = get_engine(credential_name=env.db.default_credential)
 
 
-def upload():
+def upload(disease: str, data_path: str = PYSUS_DATA_PATH):
     """
-    Connects to the EGH SQL server and load all the chunks for all
-    diseases found at `$PYSUS_DATA_PATH` into database. This method cleans
-    the chunks left.
-
+    Connects to the EpiGraphHub SQL server and load parquet chunks within
+    directories, extracted using `extract.download`, into database. Receives
+    a disease and look for local parquets paths in PYSUS_DATA_PATH, extract theirs
+    DataFrames and upsert rows to Postgres connection following EGH table
+    convention, see more in EGH's documentation:
+    https://epigraphhub.readthedocs.io/en/latest/instruction_name_tables.html#about-metadata-tables
     """
-    diseases_dir = Path(PYSUS_DATA_PATH).glob("*")
-    di_years_dir = [x for x in diseases_dir if x.is_dir()]
+    disease_years = Disease(disease).get_years(stage='all')
 
-    for dir in di_years_dir:
-        if "parquet" in Path(dir).suffix and any(os.listdir(dir)):
-            df = to_df(str(dir), clean_after_read=False)
+    for year in disease_years:
+        df = SINAN.parquets_to_df(disease, year, data_path)
+        if not df.empty:
             df.columns = df.columns.str.lower()
             df.index.name = "index"
 
-            table_i = str(dir).split("/")[-1].split(".parquet")[0]
-            table = table_i[:-4].lower()
+            tablename = "sinan_" + normalize_str(disease) + "_m"
             schema = "brasil"
 
+            print(f"Inserting {disease}-{year} on {schema}.{tablename}")
+
             with engine.connect() as conn:
                 try:
-
                     upsert(
                         con=conn,
                         df=df,
-                        table_name=table,
+                        table_name=tablename,
                         schema=schema,
                         if_row_exists="update",
                         chunksize=1000,
                         add_new_columns=True,
                         create_table=True,
                     )
 
-                    logger.info(f"Table {table} updated")
+                    print(f"Table {tablename} updated")
 
                 except Exception as e:
-                    logger.error(f"Not able to upsert {table} \n{e}")
+                    logger.error(f"Not able to upsert {tablename} \n{e}")
                     raise e
+        else:
+            print(f'[WARNING] No data for {disease} and year {year}. Skipping')
+            continue

epigraphhub/data/brasil/sinan/viz.py

@@ -1,19 +1,19 @@
 import pandas as pd
 from loguru import logger
-from pysus.online_data import SINAN
-from pysus.online_data import parquets_to_dataframe as to_df
+from pysus import SINAN
 
 from epigraphhub.connection import get_engine
 from epigraphhub.data._config import SINAN_LOG_PATH
 from epigraphhub.settings import env
 
+from . import normalize_str
+
 logger.add(SINAN_LOG_PATH, retention="7 days")
 
 engine = get_engine(credential_name=env.db.default_credential)
-aggrs = SINAN.agravos
 
 
-def parquet(ppath: str, clean_after_read=False) -> pd.DataFrame:
+def parquet(disease: str, year: str|int) -> pd.DataFrame:
     """
     Convert the parquet files into a pandas DataFrame.
 
@@ -27,14 +27,13 @@ def parquet(ppath: str, clean_after_read=False) -> pd.DataFrame:
         df (DataFrame)         : A Pandas DataFrame.
     """
 
-    df = to_df(str(ppath), clean_after_read)
-    logger.info("Parquet files converted to dataFrame")
+    df = SINAN.parquet_to_df(disease, year)
     df.columns = df.columns.str.lower()
 
     return df
 
 
-def table(disease: str, year: int) -> pd.DataFrame:
+def table(disease: str) -> pd.DataFrame:
     """
     Connect to EGH SQL server and retrieve the data by disease and year.
 
@@ -48,12 +47,10 @@ def table(disease: str, year: int) -> pd.DataFrame:
 
     """
 
-    year = str(year)[-2:].zfill(2)
-    disease = SINAN.check_case(disease)
-    dis_code = aggrs[disease].lower()
-    tablename = f"{dis_code}{year}"
+    tablename = "sinan_" + normalize_str(disease) + "_m"
+    schema = "brasil"
 
     with engine.connect() as conn:
-        df = pd.read_sql(f"SELECT * FROM brasil.{tablename}", conn)
+        df = pd.read_sql(f"SELECT * FROM {schema}.{tablename}", conn)
 
     return df