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)