from airflow.providers.postgres.hooks.postgres import PostgresHook
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
import re
[docs]
class CSVToPostgresOperator(BaseOperator):
"""
Operator to read a CSV file and save it in a Postgres database
:param csv_path: path to CSV file. For S3 use ``s3://<bucket>/<prefix>/<file>`` format.
:type csv_path: str
:param postgres_conn_id: the Airflow connection ID to use to store data in Postgres.
:type postgres_conn_id: str
:param schema_name: the Postgres schema name
:type schema_name: str
:param table_name: the Postgres table name within the schema
:type table_name: str
:param file_encoding: default to ``utf-8``. Use ``ISO-8859-1`` if CSV file contains chars outside of ``utf-8``.
:type file_encoding: str
:param if_table_exists: default to ``fail`` for data safety. Can also be ``replace`` or ``append``.
:type if_table_exists: str
:param header: how to handle the first row of the CSV. Default ``infer``. Can also be ``None``.
:type header: str
:param skip_row_count: number of rows directly after the header row (if any) to skip. Default ``0``. Useful when there are 2 or more header rows.
:type skip_row_count: int
:param compression: the compression scheme of the CSV file. Default ``infer``. Can also be ``gzip``, ``bz2``, ``zip``, ``xz``, ``None``
:type str
:param column_rename: map of ``{'oldname':'newname'}``. Default ``None``.
:type dict
:param stop_col: column number to stop at, if theres a bunch of duplicate columns
:type integer
"""
ui_color = "#A6E6A6"
@apply_defaults
def __init__(
self,
csv_path,
postgres_conn_id,
schema_name,
table_name,
file_encoding="utf-8",
if_table_exists="fail",
header="infer",
skip_row_count=0,
compression="infer",
column_rename=None,
field_names=None,
chunk_size=0,
sep=",",
normalize=False,
sheet=None, # for easy cross compatibility with xlsx operator
index_col=None,
engine=None,
stop_col=0,
*args,
**kwargs
):
super(CSVToPostgresOperator, self).__init__(*args, **kwargs)
self.csv_path = csv_path
self.file_encoding = file_encoding
self.postgres_conn_id = postgres_conn_id
self.schema_name = schema_name
self.table_name = table_name
self.if_table_exists = if_table_exists
self.header = header
self.skip_row_count = skip_row_count
self.compression = compression
self.column_rename = column_rename
self.field_names = field_names
self.chunk_size = chunk_size
self.sep = sep
self.normalize = normalize
self.index_col = index_col
self.engine = engine
self.stop_col = stop_col
[docs]
def check_table(self, dest_conn, dest_schema_table):
"""delete all contents of the given table before writing if necessary"""
result_exist = dest_conn.execute(
"SELECT to_regclass('{}');".format(dest_schema_table)
)
if result_exist.fetchone()[0]:
# wrapping pandas functionality for chunking edge cases
if self.if_table_exists == "replace":
self.log.info("Replacing existing table: {}".format(dest_schema_table))
dest_conn.execute("DROP TABLE {};".format(dest_schema_table))
elif self.if_table_exists == "fail":
raise
[docs]
def handle_df_transform(self, df, conn, if_table_exists):
"""helper to write csv"""
if self.field_names:
df.columns = self.field_names
if self.column_rename is not None:
df.rename(mapper=self.column_rename, axis="columns", inplace=True)
if self.normalize:
norm_cols = {
col: re.sub(
r"\W+",
"",
str(col)
.replace(" ", "_")
.replace("-", "_")
.replace(".", "_")
.replace("/", "_")
.lower(),
)
for col in df.columns.tolist()
}
df.rename(mapper=norm_cols, axis="columns", inplace=True)
if self.stop_col:
df = df.iloc[:, : -(len(df.columns.tolist()) - self.stop_col)]
df.to_sql(
self.table_name,
conn,
schema=self.schema_name,
if_exists=if_table_exists,
method="multi",
)
[docs]
def execute(self, context):
import pandas as pd
from sqlalchemy import create_engine
pg_hook = PostgresHook(postgres_conn_id=self.postgres_conn_id)
engine = create_engine(pg_hook.get_uri())
conn = engine.connect()
if self.chunk_size:
self.check_table(conn, "{}.{}".format(self.schema_name, self.table_name))
for df in pd.read_csv(
self.csv_path,
iterator=True,
chunksize=self.chunk_size,
sep=self.sep,
header=self.header,
):
self.handle_df_transform(df, conn, "append")
else:
df = pd.read_csv(
self.csv_path,
encoding=self.file_encoding,
header=self.header,
compression=self.compression,
sep=self.sep,
engine=self.engine,
index_col=self.index_col,
)
if self.skip_row_count > 0:
while self.skip_row_count > 0:
df.drop(0, inplace=True)
self.skip_row_count -= 1
df = df.infer_objects()
self.handle_df_transform(df, conn, self.if_table_exists)
conn.close()