from prefect import flow, task, get_run_logger
import pandas as pd
import boto3
import requests
import io

@task(name="Getting data")
def extract(return_df=True):
    if return_df:
        url = 'https://www.federalreserve.gov/datadownload/Output.aspx?rel=H8&series=dd48166f12d986aede821fb86d9185d7&lastobs=&from=&to=&filetype=csv&label=include&layout=seriescolumn&type=package'
        r = requests.get(url)
        df = pd.read_csv(io.StringIO(r.content.decode('utf-8')))
        return df
    else:
        return 'this is a string'

@task(name="Doing something with data")
def transform(df_or_string):
    if isinstance(df_or_string, str):
        print('this is a string')
        the_string = df_or_string +' more string'
        return the_string
    else:
        df = df_or_string
        df['new_col'] = 'Add this constant'
        return df

@task(name="Log output of transform")
def load(df):
    logger = get_run_logger()
    logger.info(df)
    
@flow(name="Retry Test")
def main_flow():
    data_frame = extract(return_df = True)  ### caches on retry
    df_transformed = transform(data_frame) ### DOES NOT CACHE ON RETRY BECAUSE IT HAS AN INPUT THAT IS A DATAFRAME TYPE CACHED TASK RESULT
    _ = load(df_transformed) ### DOES NOT CACHE ON RETRY BECAUSE IT HAS AN INPUT THAT IS A DATAFRAME TYPE CACHED TASK RESULT

    string_text = extract(return_df = False) ### caches on retry
    longer_string = transform(string_text) ### caches on retry
    _ = load(longer_string) ### caches on retry

    data_frame2 = extract(return_df = True)  ### this caches on initial run and on caches on retry
    df_transformed2 = transform(data_frame2) ### DOES NOT CACHE ON INTITIAL RUN OR RETRY BECAUSE IT HAS AN INPUT THAT IS A DATAFRAME TYPE CACHED TASK RESULT
    _ = load(df_transformed2) ### THIS CACHES ON INTIIAL RUN AND ON RETRY..... not sure why this is different than the line above...???????

if __name__ == "__main__":
    main_flow()