Notes

from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
from typing import Dict
__dag__ = "create-upcoming-years-bq-tables"
from hrastro.plugins.utils import CommonUtil
from hrastro.plugins.utils import appConstants as AC
from hrastro.plugins.utils import stringConstants as SC
from airflow.providers.google.cloud.transfers.gcs_to_gcs import GCSToGCSOperator
from hrastro.plugins.operators.CreateDataProcCluster import CreateDataProcCluster
from hrastro.plugins.operators.DataprocSparkOperator import DataprocSparkOperator
from hrastro.plugins.operators.DeleteDataProcCluster import DeleteDataProcCluster
from airflow import AirflowException
from google.cloud import bigquery
from datetime import datetime, timedelta
from airflow.operators.python_operator import PythonOperator
from airflow import DAG
import os



DEFAULT_ARG: Dict = {
SC.GCP_CONN_ID: 'gcp_etl_conn_id'
,'depends_on_past': False
, SC.PROJECT_ID: CreateDataProcCluster.get_project_id_from_connection(gcp_conn_id='gcp_etl_conn_id')
, SC.REGION: CreateDataProcCluster.get_dataproc_region()
, SC.DOMAIN_NAME: AC.WS_VAR.get(SC.DOMAIN_NAME)
, SC.OWNER: AC.WS_VAR.get(SC.OWNER)
, SC.RUN_BOOK: ''
, SC.TR_PRODUCT_ID: AC.WS_VAR.get(SC.TR_PRODUCT_ID)
, SC.DAG_FILE_NAME: os.path.splitext(os.path.basename(__file__))[0]
, SC.APPLICATION: 'example'
, SC.PRIORITY: 'P1'
, SC.PRODUCT_LINE: "example",
"owner": "DSI-HR",
"description": "To create upcoming year tables in BQ",
"depend_on_past": False,
"start_date": datetime(2020, 5, 2),
"email_on_failure": False,
"email_on_retry": False,
"retries": 1
}






dag = DAG(
dag_id='create-upcoming-years-bq-tables',
default_args=DEFAULT_ARG,
description="To create upcoming year tables in BQ",
schedule_interval=None,
catchup=False,
dagrun_timeout=timedelta(minutes=60),
max_active_runs=1,
)



def util_bq_operation():
try:
# Construct a BigQuery client object.
bq_hook = BigQueryHook(bigquery_conn_id="gcp_etl_conn_id",use_legacy_sql=False)
client = bigquery.Client(project=bq_hook._get_field("project"),credentials=bq_hook._get_credentials())

# replace 2022, with 2023 and change where condition to have 2022 tables
sql = """SELECT table_name,replace(replace(ddl,"2022","2023"),"CREATE TABLE","CREATE TABLE IF NOT EXISTS")ddl FROM
WW_HR_DL_US_PAYROLL_REPORT_TABLES.INFORMATION_SCHEMA.TABLES
where table_name like '%2022'"""

# Start the query, passing in the extra configuration.
query_job = client.query(sql) # Make an API request.
response_bq = query_job.result() # Wait for the job to complete.
for row in response_bq:
upcoming_year_ddl = row.ddl
this_year_table = row.table_name
create_upcoming_year_tbl = client.query(upcoming_year_ddl)
create_upcoming_year_tbl.result()
print(f"Upcoming year table for this year table"
f" {this_year_table} "
f"has been created successfully")

except Exception as e:
raise AirflowException(
'Error while creating upcoming year tables in BQ'.format(
str(e)))

util_bq_ddl_operation_task = PythonOperator(
task_id='util_bq_ddl_operation',
default_args=DEFAULT_ARG,
python_callable=util_bq_operation,
dag=dag
)
util_bq_ddl_operation_task