diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 0000000..7458515 --- /dev/null +++ b/.editorconfig @@ -0,0 +1,17 @@ +# EditorConfig file: http://EditorConfig.org +root = true + +[*] +charset = utf-8 +end_of_line = lf +indent_size = 4 +indent_style = space +insert_final_newline = true +trim_trailing_whitespace = true + +[*.py] +indent_size = 4 +max_line_length = 88 + +[*.md] +trim_trailing_whitespace = false diff --git a/.gitignore b/.gitignore index 061db0f..470bdc1 100644 --- a/.gitignore +++ b/.gitignore @@ -213,4 +213,5 @@ __marimo__/ *.ipynb /data-dir /uploads -*.DS_Store \ No newline at end of file +*.DS_Store +*.pkl \ No newline at end of file diff --git a/config.yaml b/config.yaml index e71fcaf..b322502 100644 --- a/config.yaml +++ b/config.yaml @@ -9,8 +9,13 @@ status_app: bot_name: "snt-maxxer" channels: - https://status.app/c/G6EAAMSs5eYUrSjkDriqGHx1OITK3bd8aUlQKA9M5Mg08uTbwYKNMVxLXxDGfzde3Ub9OeDeNCmVTbP-vZs-rsCtWIUKDBBWUBXrEaGpJQ5Kaj0o4pYlcJ0iLlnP-MQRxCRwy3pE3JOiMxYYIyb4WtmaksZHTHQKCLUc14iNpWoidEbVeeO2g931cXu8Lns3Bw==#zQ3shsFYujbDQdRhSKS9RHuGCwxHQ1WLkNYvGPRksf4ebDWFW - + backend_params: url: "http://localhost:8080" logLevel: "INFO" - data_dir: "./data-dir" \ No newline at end of file + data_dir: "./data-dir" + +sleep: + # Values must be in MINUTES + upload: 15 # 15 min + download: 1440 # 1 day diff --git a/download.py b/download.py index c7de877..baf685b 100644 --- a/download.py +++ b/download.py @@ -1,31 +1,72 @@ from clients.status_backend import StatusBackend +from typing import Optional, Any +import pandas as pd import constants, data_utils -import os, datetime, pickle +import os, datetime, pickle, time, json, logging -if __name__ == "__main__": +def save_pkl(file_path: str, data: Any): + """ + Save the given data into a Pickel file + Parameters: + - `file_path` - the path where the data will be saved + - `data` - the data that will be saved + """ + with open(file_path, "wb") as f: + pickle.dump(data, f) + +def run(logger: logging.Logger): backend = StatusBackend(**constants.STATUS_BACKEND_PARAMS) info = data_utils.login(backend, constants.CONFIG["status_app"]["bot_name"]) - + logger.info(f"Logged in with account {constants.CONFIG['status_app']['bot_name']}") + message_folder = os.path.join(constants.UPLOAD_PATH, "messages") community_folder = os.path.join(constants.UPLOAD_PATH, "channel_info") + latest_batch_path = os.path.join(constants.UPLOAD_PATH, "start_timestamp.pkl") + start_timestamp: Optional[datetime.datetime] = pd.read_pickle(latest_batch_path) if os.path.exists(latest_batch_path) else None + for channel_url in constants.CONFIG["status_app"]["channels"]: + logger.info(f"Starting {channel_url}") community = data_utils.get_community_info(backend, channel_url) current_community_folder = os.path.join(community_folder, community["community_id"]) os.makedirs(current_community_folder, exist_ok=True) file_path = os.path.join(current_community_folder, str(datetime.datetime.now().timestamp()).replace(".", "") + ".pkl") - - with open(file_path, "wb") as f: - pickle.dump(community, f) + save_pkl(file_path, community) for channel in community["channels"]: + msg = f"Downloading messages from {community['community_name']} #{channel['channel_name']}" + if start_timestamp: + msg += f" (from {start_timestamp} onwards)" + logger.info(msg) params = { "backend": backend, "chat_id": channel["chat_id"], "folder": message_folder, - "community_info": channel + "community_info": channel, + "start_timestamp": start_timestamp } - data_utils.save_messages(backend, channel["chat_id"], message_folder, channel) + data_utils.save_messages(**params) - backend.logout() \ No newline at end of file + backend.logout() + logger.info(f"Logged out of {constants.CONFIG['status_app']['bot_name']}") + + timestamps = [] + for file_name in os.listdir(message_folder): + with open(os.path.join(message_folder, file_name), "r") as f: + data: dict = json.load(f) + timestamps.append(data["metadata"]["latest_msg_timestamp"]) + + if not timestamps: + return + + save_pkl(latest_batch_path, datetime.datetime.fromtimestamp(max(timestamps))) + +if __name__ == "__main__": + + logger = data_utils.get_logger("download") + while True: + run(logger) + seconds = 60 * constants.CONFIG["sleep"]["download"] + logger.info(f"Sleeping for {constants.CONFIG['sleep']['download']} minutes") + time.sleep(seconds) diff --git a/upload.py b/upload.py index 0c737c9..e0906b3 100644 --- a/upload.py +++ b/upload.py @@ -1,5 +1,5 @@ -import constants -import json, datetime, os +import constants, data_utils +import json, datetime, os, time, logging import pandas as pd from pathlib import Path, PosixPath from postgres import Postgres @@ -16,6 +16,8 @@ def get_community_members(file_path: PosixPath) -> pd.DataFrame: - processed data to be uploaded to the database """ info: dict = pd.read_pickle(file_path) + if not isinstance(info, dict): + return pd.DataFrame() data = pd.DataFrame(info["members"]["info"]).assign( community_total = info["members"]["total"], @@ -30,14 +32,17 @@ def get_community_info(file_path: PosixPath) -> pd.DataFrame: Parameters: - `file_path` - the `.pkl` file that has overall member information from `data_utils.get_community_info` - + Output: - Single row for the given file. The data is returned in a DataFrame to make data uploading more robust """ info: dict = pd.read_pickle(file_path) # members data can be found in function get_community_members - info.pop("members") - return pd.DataFrame([info]) + if isinstance(info, dict): + info.pop("members") + return pd.DataFrame([info]) + + return pd.DataFrame() def get_messages(file_paths: list[str]) -> pd.DataFrame: """ @@ -49,6 +54,9 @@ def get_messages(file_paths: list[str]) -> pd.DataFrame: Output: - DataFrame with all of the messages """ + if not file_paths: + return pd.DataFrame() + raw_data = [] for file_path in file_paths: with open(file_path, "r") as f: @@ -61,11 +69,10 @@ def get_messages(file_paths: list[str]) -> pd.DataFrame: if not column.endswith("timestamp"): continue data[column] = pd.to_datetime(data[column], unit="s") - + return data.copy() - -def upload(data: dict[str, pd.DataFrame], connector: Postgres): +def upload(data: dict[str, pd.DataFrame], connector: Postgres, logger: logging.Logger): """ Upload the raw data to Postgres. @@ -73,9 +80,11 @@ def upload(data: dict[str, pd.DataFrame], connector: Postgres): - `data` - the raw concatenated data - `connector` - initialized Postgres connection """ + schema = constants.CONFIG["postgres"]["schema"] for data_key, table_name in constants.CONFIG["postgres"]["tables"].items(): df = data[data_key] if len(df) == 0: + logger.info(f"No data to upload to {schema}.{table_name}") continue df = df.assign( @@ -87,31 +96,27 @@ def upload(data: dict[str, pd.DataFrame], connector: Postgres): for column in df.columns if isinstance(df[column].dropna().reset_index(drop=True).iloc[0], (dict, list)) ] - connector.insert(df, table_name, constants.CONFIG["postgres"]["schema"], json_columns) + connector.insert(df, table_name, schema, json_columns) + logger.info(f"Uploaded {len(df)} row(s) to {schema}.{table_name}") -if __name__ == "__main__": +def run(username: str, password: str, host: str, database: str, port: str, logger: logging.Logger): completed = [] path = Path(constants.UPLOAD_PATH) - load_dotenv() - connector = connector = Postgres( - username = os.getenv("POSTGRES_USERNAME"), - password = os.getenv("POSTGRES_PASSWORD"), - host = os.getenv("POSTGRES_HOST"), - database = os.getenv("POSTGRES_DATABASE"), - port = int(os.getenv("POSTGRES_PORT")) - ) + connector = Postgres(username, password, port, database, host) + logger.info(f"Initialized Postgres connector!") # NOTE: the keys are the same as in config.yaml -> postgres.tables data = { "members": [], "community": [], } - + for file_path in path.rglob("*.pkl"): data["community"].append(get_community_info(file_path)) data["members"].append(get_community_members(file_path)) completed.append(file_path) - + logger.info(f"Created community and member data for {file_path}") + for key, value in data.items(): data[key] = pd.concat(value, ignore_index=True) @@ -119,9 +124,31 @@ if __name__ == "__main__": data["messages"] = get_messages(file_paths) completed += file_paths - upload(data, connector) + upload(data, connector, logger) connector.close() - if completed: - for file_path in completed: - os.remove(file_path) \ No newline at end of file + if not completed: + return + + for file_path in completed: + os.remove(file_path) + logger.info(f"Deleted {file_path}") + +if __name__ == "__main__": + + logger = data_utils.get_logger("upload") + load_dotenv() + params = { + "username": os.getenv("POSTGRES_USERNAME"), + "password": os.getenv("POSTGRES_PASSWORD"), + "host": os.getenv("POSTGRES_HOST"), + "database": os.getenv("POSTGRES_DATABASE"), + "port": int(os.getenv("POSTGRES_PORT")), + "logger": logger + } + + while True: + run(**params) + seconds = 60 * constants.CONFIG["sleep"]["upload"] + logger.info(f"Sleeping for {constants.CONFIG['sleep']['upload']} minutes") + time.sleep(seconds)