Fixed: TSV header getting included and passed to postgres table

Author: stemitom

dags/operators/csv_to_postgres.py

@@ -4,7 +4,7 @@
 
 class LoadCsvtoPostgresOperator(BaseOperator):
     """
-    Moves data from a comma seperated file to Postgres
+    Moves data from a delimiter(tab) seperated file to Postgres
     """
     template_fields = ("file_path",)
 
@@ -22,6 +22,6 @@ def execute(self, context):
             postgres.bulk_load(self.table, self.file_path)
         except Exception as err:
             self.log.error(err)
-            return err
+            raise ValueError(err)
         else:
             self.log.info(f"Loaded file {self.file_path} into table {self.table}")

dags/pipeline.py

@@ -46,8 +46,9 @@ def _fetch_data(outfile):
 def _transform_to_csv(infile, outfile):
     pathlib.Path("/tmp/data/stg").mkdir(parents=True, exist_ok=True)
     data = pd.read_json(infile)
-    data = data.set_index("date_of_interest")
-    data.to_csv(outfile)
+    data['date_of_interest'] = data['date_of_interest']\
+        .apply(lambda date: date.strip('T00:00:00.000'))
+    data.to_csv(outfile, index=False, header=False, sep="\t", encoding="utf-8")
     logging.info(f"INFO: Processed {infile} and moved it to {outfile}")
 
 
@@ -66,7 +67,7 @@ def _transform_to_csv(infile, outfile):
     dag=dag,
     op_kwargs={
         "infile": "/tmp/data/raw/covid_data_{{ ds }}.json",
-        "outfile": "/tmp/data/raw/covid_data_{{ ds }}.csv",
+        "outfile": "/tmp/data/stg/covid_data_{{ ds }}.csv",
     },
 )
 
@@ -83,53 +84,17 @@ def _transform_to_csv(infile, outfile):
 create_covid_data_table = PostgresOperator(
     task_id="create_table_covid",
     postgres_conn_id="covid_postgres",
-    sql="sql/create_table.sql",
+    sql="sql/create_test.sql",
     dag=dag,
 )
 
 load_csv_to_postgres_dwh = LoadCsvtoPostgresOperator(
     task_id='load_to_covid_data_table',
     postgres_conn_id="covid_postgres",
-    table="covid_data",
+    table="covid_table",
     file_path="/tmp/data/stg/covid_data_{{ ds }}.csv",
     dag=dag,
 )
 
 # fetch_data >> transform_to_csv >> normalize_covid_csv >> create_covid_data_table>> load_csv_to_postgres_dwh
-
-def _test_transform(infile, outfile):
-    pathlib.Path("/tmp/data/stg").mkdir(parents=True, exist_ok=True)
-    data = pd.read_json(infile)
-    data = data.set_index("date_of_interest")
-    data = data[['date_of_interest', 'case_count']]
-    data.to_csv(outfile)
-    logging.info(f"INFO: Processed {infile} and moved it to {outfile}")
-
-test_transform = PythonOperator(
-    task_id="test_transform",
-    python_callable=_test_transform,
-    dag=dag,
-    op_kwargs={
-        "infile": "/tmp/data/raw/covid_data_{{ ds }}.json",
-        "outfile": "/tmp/data/raw/covid_data_{{ ds }}.csv",
-    },
-
-)
-
-create_covid_test_table = PostgresOperator(
-    task_id="create_covid_test_table",
-    postgres_conn_id="covid_postgres",
-    sql="sql/create_test.sql",
-    dag=dag
-)
-
-
-test_data_load = LoadCsvtoPostgresOperator(
-    task_id='test_data_load',
-    postgres_conn_id="covid_postgres",
-    table="covid_test",
-    file_path="/tmp/data/stg/covid_data_{{ ds }}.csv",
-    dag=dag,
-)
-
-fetch_data >> test_transform >> normalize_covid_csv >> create_covid_test_table >> test_data_load
+fetch_data >> transform_to_csv >> create_covid_data_table >> load_csv_to_postgres_dwh

dags/sql/create_table.sql

@@ -1,7 +1,7 @@
 DROP TABLE IF EXISTS covid_data;
 
 CREATE TABLE IF NOT EXISTS covid_data (
-    date DATE PRIMARY KEY,
+    date_of_interest DATE PRIMARY KEY,
     case_count INT,
     probable_case_count INT,
     hospitalized_count INT,