diff --git a/pypistats/tasks/pypi.py b/pypistats/tasks/pypi.py index 55d317f..2b3dedc 100644 --- a/pypistats/tasks/pypi.py +++ b/pypistats/tasks/pypi.py @@ -69,9 +69,10 @@ def get_google_credentials(): return credentials -def get_daily_download_stats(env="dev", date="None"): +def get_daily_download_stats(event, context): """Get daily download stats for pypi packages from BigQuery.""" start = time.time() + env = event["env"] if os.environ.get("ENV", None) is None: load_env_vars(env) job_config = bigquery.QueryJobConfig() @@ -80,7 +81,8 @@ def get_daily_download_stats(env="dev", date="None"): project=os.environ["GOOGLE_PROJECT_ID"], credentials=credentials ) - if date == "None": + date = event.get("date", None) + if date is None: date = str(datetime.date.today() - datetime.timedelta(days=1)) # # Prepare a reference to the new dataset @@ -155,8 +157,7 @@ def get_daily_download_stats(env="dev", date="None"): def update_db(df, env="dev"): """Update the db with new data by table.""" - if os.environ.get("ENV", None) is None: - load_env_vars(env) + load_env_vars(env) connection, cursor = get_connection_cursor(env) df_groups = df.groupby("category_label") @@ -204,16 +205,19 @@ def update_table(connection, cursor, table, df, date): return False -def update_all_package_stats(env="dev", date="None"): +def update_all_package_stats(event, context): """Update stats for __all__ packages.""" print("__all__") start = time.time() - if date == "None": + date = event.get("date", None) + if date is None: date = str(datetime.date.today() - datetime.timedelta(days=1)) + env = event["env"] if os.environ.get("ENV", None) is None: load_env_vars(env) + connection, cursor = get_connection_cursor(env) success = {} @@ -245,16 +249,19 @@ def update_all_package_stats(env="dev", date="None"): return success -def update_recent_stats(env="dev", date="None"): +def update_recent_stats(event, context): """Update daily, weekly, monthly stats for all packages.""" print("recent") start = time.time() - if date == "None": + date = event.get("date", None) + if date is None: date = str(datetime.date.today() - datetime.timedelta(days=1)) + env = event["env"] if os.environ.get("ENV", None) is None: load_env_vars(env) + connection, cursor = get_connection_cursor(env) downloads_table = "overall" @@ -315,16 +322,20 @@ def get_connection_cursor(env): return connection, cursor -def purge_old_data(env="dev", age=MAX_RECORD_AGE, date="None"): +def purge_old_data(event, context): """Purge old data records.""" print("Purge") + age = MAX_RECORD_AGE start = time.time() - if date == "None": + date = event.get("date", None) + if date is None: date = str(datetime.date.today() - datetime.timedelta(days=1)) + env = event["env"] if os.environ.get("ENV", None) is None: load_env_vars(env) + connection, cursor = get_connection_cursor(env) date = datetime.datetime.strptime(date, '%Y-%m-%d') @@ -439,9 +450,14 @@ def get_query(date): if __name__ == "__main__": - date = "2018-04-16" + date = "2018-04-17" env = "prod" + event = { + "date": date, + "env": env, + } + context = None print(date, env) - print(get_daily_download_stats(env, date)) - print(update_all_package_stats(env, date)) - print(update_recent_stats(env, date)) + print(get_daily_download_stats(event, context)) + print(update_all_package_stats(event, context)) + print(update_recent_stats(event, context)) diff --git a/zappa_settings_template.json b/zappa_settings_template.json index 24783e2..5d18627 100644 --- a/zappa_settings_template.json +++ b/zappa_settings_template.json @@ -11,22 +11,22 @@ "events": [ { "function": "pypistats.tasks.pypi.get_daily_download_stats", - "kwargs": {"env": "dev", "date": "None"}, + "kwargs": {"env": "dev"}, "expression": "cron(0 1 * * ? *)" }, { "function": "pypistats.tasks.pypi.update_all_package_stats", - "kwargs": {"env": "dev", "date": "None"}, + "kwargs": {"env": "dev"}, "expression": "cron(10 1 * * ? *)" }, { "function": "pypistats.tasks.pypi.update_recent_stats", - "kwargs": {"env": "dev", "date": "None"}, + "kwargs": {"env": "dev"}, "expression": "cron(15 1 * * ? *)" }, { "function": "pypistats.tasks.pypi.purge_old_data", - "kwargs": {"env": "dev", "date": "None"}, + "kwargs": {"env": "dev"}, "expression": "cron(20 1 * * ? *)" } ], @@ -69,22 +69,22 @@ "events": [ { "function": "pypistats.tasks.pypi.get_daily_download_stats", - "kwargs": {"env": "prod", "date": "None"}, + "kwargs": {"env": "prod"}, "expression": "cron(0 1 * * ? *)" }, { "function": "pypistats.tasks.pypi.update_all_package_stats", - "kwargs": {"env": "prod", "date": "None"}, + "kwargs": {"env": "prod"}, "expression": "cron(10 1 * * ? *)" }, { "function": "pypistats.tasks.pypi.update_recent_stats", - "kwargs": {"env": "prod", "date": "None"}, + "kwargs": {"env": "prod"}, "expression": "cron(15 1 * * ? *)" }, { "function": "pypistats.tasks.pypi.purge_old_data", - "kwargs": {"env": "prod", "date": "None"}, + "kwargs": {"env": "prod"}, "expression": "cron(20 1 * * ? *)" } ],