def execute(self, context):
snowflake_hook = SnowflakeHook(
snowflake_conn_id=self.snowflake_conn_id)
# Snowflake won't accept list of files it has to be tuple only.
# but in python tuple([1]) = (1,) => which is invalid for snowflake
files = str(self.s3_keys)
files = files.replace('[', '(')
files = files.replace(']', ')')
# we can extend this based on stage
base_sql = """
FROM @{stage}/
files={files}
file_format={file_format}
""".format(stage=self.stage,
files=files,
file_format=self.file_format)
if self.columns_array:
copy_query = """
COPY INTO {schema}.{table}({columns}) {base_sql}
""".format(schema=self.schema,
table=self.table,
columns=",".join(self.columns_array),
base_sql=base_sql)
else:
copy_query = """
COPY INTO {schema}.{table} {base_sql}
""".format(schema=self.schema, table=self.table, base_sql=base_sql)
self.log.info('Executing COPY command...')
print(snowflake_hook.get_uri())
snowflake_hook.run(copy_query, self.autocommit)
self.log.info("COPY command completed")
def get_table_max_incremental_number(**context):
get_table_max_incremental_number = MEDICARE_PART_D_PROVIDER_TABLE_MAX_NPI
print("Get Table MAX increment number...")
dwh_hook = SnowflakeHook(snowflake_conn_id="kyu_snowflake_conn",
warehouse="COMPUTE_WH")
dwh_hook.run(get_table_max_incremental_number, autocommit=True)
table_max_npi = dwh_hook.get_first(get_table_max_incremental_number)[0]
print("MAX Table increment number: " + str(table_max_npi))
def execute(self, context):
hook = SnowflakeHook(snowflake_conn_id=snowflake_conn_id)
with snowflake_query_tracker(database=database) as st:
hook.run(select_query)
session_id, query_id = st.get_last_session_with_query_id(
many=False)
context["ti"].xcom_push(key="session_id", value=session_id)
context["ti"].xcom_push(key="query_id", value=query_id)
def create_or_replace_view(**context):
create_or_replace_view = MEDICARE_PART_D_PROVIDER_LOAD_CREATE_VIEW
print("creating or replacing existing view.")
dwh_hook = SnowflakeHook(snowflake_conn_id="kyu_snowflake_conn",
warehouse="COMPUTE_WH")
dwh_hook.run(create_or_replace_view, autocommit=True)
def load_stage_to_master(**context):
load_stage_to_master = MEDICARE_PART_D_PROVIDER_LOAD_MASTER_TABLE
print("Loading stage to master...")
dwh_hook = SnowflakeHook(snowflake_conn_id="kyu_snowflake_conn",
warehouse="COMPUTE_WH")
dwh_hook.run(load_stage_to_master, autocommit=True)
def stage_file(**context):
stage_file = MEDICARE_PART_D_PROVIDER_COPY_TO_STAGE
print("Loading S3 file into staging table")
dwh_hook = SnowflakeHook(snowflake_conn_id="kyu_snowflake_conn",
warehouse="COMPUTE_WH")
dwh_hook.run(stage_file, autocommit=True)
def truncate_stage_table(**context):
truncate_stage_table = MEDICARE_PART_D_PROVIDER_TRUNCATE_STAGE
print("Truncate staging table")
dwh_hook = SnowflakeHook(snowflake_conn_id="kyu_snowflake_conn",
warehouse="COMPUTE_WH")
dwh_hook.run(truncate_stage_table, autocommit=True)
def update_customers(**kwargs):
snowflake_hook = SnowflakeHook(snowflake_conn_id=SNOWFLAKE_CONNECTION_ID)
snowflake_hook.run(update_query)
def execute(self, context):
hook = SnowflakeHook(snowflake_conn_id=snowflake_conn_id)
with snowflake_query_tracker(database=database):
hook.run(sql=select_query)
def update_customers_with_monitoring(**kwargs):
snowflake_hook = SnowflakeHook(snowflake_conn_id=SNOWFLAKE_CONNECTION_ID)
with snowflake_query_tracker(log_tables=False, database=database):
snowflake_hook.run(update_query)
