
# Conection and setup Run the following in Airflow webserver container
'''
CREATE TABLE pageview_counts (pagename VARCHAR(50) NOT NULL, pageviewcount INT NOT NULL, datetime TIMESTAMP NOT NULL);

To setup connection between a DAG and postgres database either use a CLI or web interface to add connection.
For simplicity I've used user interface
Click on Admin-> Connections-> Add Connection->
Connection Id: postgres_example
Connection Type: postgres
Host: postgres (This is because we are using docker compose and host url is the service name for database)
Login: airflow
Password: <password for the postgres database used for airflow> **DO NOT DO THIS IN PRODUCTION **
->Save<-
'''

from urllib import request
import airflow
from airflow import DAG
from airflow.operators.bash import BashOperator
from airflow.operators.python import PythonOperator
from airflow.providers.postgres.operators.postgres import PostgresOperator

# Create a dag to get wikipedia data, extract results, store sql statements on a .sql file and set up a hook to postgres
# database

# DAG to download wikipedia details hourly
dag = DAG(
    dag_id="stocksense",
    start_date=airflow.utils.dates.days_ago(1),
    schedule_interval="@hourly",
    template_searchpath="/tmp"
)
# Task
def _get_data(execution_date):
    year, month, day, hour, *_ = execution_date.timetuple()
    url = ("https://dumps.wikimedia.org/other/pageviews/" 
           f"{year}/{year}-{month:0>2}/" 
           f"pageviews-{year}{month:0>2}02-{hour:0>2}0000.gz"
    )
    print(f"URL is {url}")
    output_path = "/tmp/wikipageviews.gz"
    request.urlretrieve(url, output_path)

#Task
get_data = PythonOperator(task_id="get_data",
                          python_callable=_get_data,
                          dag=dag,
)

#Task
extract_gz = BashOperator(
    task_id="extract_gz",
    bash_command="gunzip --force /tmp/wikipageviews.gz", dag=dag,
)


# Task to fetch page views passed in as **kwargs
def _fetch_pageviews(pagenames, execution_date, **_):
    result = dict.fromkeys(pagenames, 0)
    with open("/tmp/wikipageviews", "r") as f:
       for line in f:
           domain_code, page_title, view_counts, _ = line.split(" ")
           if domain_code == "en" and page_title in pagenames:
               result[page_title] = view_counts
               # print(result)
               # Prints e.g. "{'Facebook': '778', 'Apple': '20', 'Google': '451', 'Amazon': '9', 'Microsoft': '119'}"

    with open("/tmp/postgres_query.sql", "w") as f:
        for pagename, pageviewcount in result.items():
            f.write(
                "INSERT INTO pageview_counts VALUES (" f"'{pagename}', {pageviewcount}, '{execution_date}'" ");\n"
            )

fetch_pageviews = PythonOperator(
    task_id="fetch_pageviews",
    python_callable=_fetch_pageviews,
    op_kwargs={"pagenames":
                   {
                        "Google",
                        "Amazon",
                        "Apple",
                        "Microsoft",
                        "Facebook"
                    }
               },
    dag=dag,
)

# Task to write to postgres database from a sql file that was generated by fetch_pageviews
# For this example use default postgres connection
write_to_postgres = PostgresOperator(
    task_id="write_to_postgres",
    postgres_conn_id="postgres_example",
    sql="postgres_query.sql",
    dag=dag,
)

get_data >> extract_gz >> fetch_pageviews >> write_to_postgres