diff --git a/.dockerignore b/.dockerignore deleted file mode 100644 index 256108d..0000000 --- a/.dockerignore +++ /dev/null @@ -1,14 +0,0 @@ -# Git -.git -.gitignore - -# Python -__pycache__/ -*.pyc - -# Status Bot -uploads/ -bot/docs - -config.yaml -.env diff --git a/Dockerfile b/Dockerfile deleted file mode 100644 index c406581..0000000 --- a/Dockerfile +++ /dev/null @@ -1,13 +0,0 @@ -FROM python:3.12-slim - -WORKDIR /app -COPY requirements.txt . -COPY bot/requirements.txt bot/requirements.txt - -RUN pip install --no-cache-dir -r requirements.txt \ - && if [ -f bot/requirements.txt ]; then pip install --no-cache-dir -r bot/requirements.txt; fi - -COPY . . - -ENTRYPOINT ["python", "monitor.py"] -CMD [] diff --git a/Jenkinsfile b/Jenkinsfile deleted file mode 100644 index 66a95e9..0000000 --- a/Jenkinsfile +++ /dev/null @@ -1,71 +0,0 @@ -#!/usr/bin/env groovy -library 'status-jenkins-lib@v1.9.16' - -pipeline { - agent { - docker { - label 'linuxcontainer' - image 'harbor.status.im/infra/ci-build-containers:linux-base-1.0.0' - args '--volume=/var/run/docker.sock:/var/run/docker.sock ' + - '--user jenkins' - } - } - - options { - disableConcurrentBuilds() - /* manage how many builds we keep */ - buildDiscarder(logRotator( - numToKeepStr: '20', - daysToKeepStr: '30', - )) - } - parameters { - string( - name: 'DOCKER_CRED', - description: 'Name of Docker Registry credential.', - defaultValue: params.DOCKER_CRED ?: 'harbor-status-im-robot', - ) - string( - name: 'DOCKER_REGISTRY_URL', - description: 'URL of the Docker Registry', - defaultValue: params.DOCKER_REGISTRY_URL ?: 'https://harbor.status.im', - ) - string( - name: 'IMAGE_TAG', - description: 'Image tag', - defaultValue: params.IMAGE_TAG ?: 'master', - ) - string( - name: 'IMAGE_NAME', - description: 'Name of the Docker image', - defaultValue: 'bi/status-bot', - ) - } - - stages { - stage('Bulding docker images') { - steps { - script { - image = docker.build( - "${params.IMAGE_NAME}:${params.IMAGE_TAG}", "./" - ) - } - } - } - stage('Push docker image'){ - steps { - script { - withDockerRegistry([ - credentialsId: params.DOCKER_CRED, url: params.DOCKER_REGISTRY_URL - ]) { - image.push() - } - } - } - } - } - - post { - cleanup { cleanWs() } - } -} diff --git a/README.md b/README.md index efbcd9c..6254086 100644 --- a/README.md +++ b/README.md @@ -1,139 +1,89 @@ -# [Status App Community Monitoring](https://status.app/) +# Status Python SDK -Monitoring tool for Status App communities. **No personal data is collected from users.** +![Status Python SDK header image](./docs/images/overview-header.png) +The initial Python Status Backend was built with testing in mind, instead of easy developer access. The objective of this repository is to make a SDK that is: -| Field | Hashed | Description | -|:----------------------|:---------|:------------------------------------------------------------| -| **id** | **Yes** | The message's ID | -| **whisper_timestamp** | No | The whisper timestamp of the message | -| **from** | **Yes** | The public key of the user | -| **message_type** | No | The message type | -| **seen** | No | True if the message has been seen otherwise False | -| **chat_id** | No | The chat ID is a combination of community ID and channel ID | -| **community_id** | No | The ID of the community | -| **response_to** | **Yes** | Ithe public key of the user who the response is for | -| **timestamp** | No | The timestamp of the message | -| **deleted** | No | True if the message was deleted otherwise False | +- **light** - as less external packages when it comes to working with Status App +- **fast** - quick to get started with Status Python +- **documented** - clear explanations of what was done and **why it was done in a specific way**. -Status Bot account information can be found in [`config.yaml`](./config.yaml). +Currently this repository is not on [PyPi](https://pypi.org/) but will be added when core functionality has been devleoped and tested. ## How it works ```mermaid -graph LR - subgraph Communities[Status App] - subgraph Status[Status Community] - StatusMessages[Messages] - StatusInfo[Information] - end - subgraph Logos[Logos Community] - LogosMessages[Messages] - LogosInfo[Information] - end +graph TB + subgraph backend[status-im/status-go] + subgraph Endpoints[Network: status-bridge] + RPC[RPC] + HTTP[REST] + SOCKET[Web Socket] + end + Vol[(Backup)] + end + + subgraph bot[Python SDK] + REQUIREMENTS[requirements.txt] + SDK[class Account] + SIGNAL[class Signal] end - subgraph Bot[Docker Container] - RawDataLocal[(Raw Data)] - Script[monitor.py] + subgraph external[External Services] + COINGECKO[CoinGecko] + EVM end - subgraph IFT[IFT Infrastructure] - RawDataIFT[(Raw Data)] - ProcessedDataIFT[(Processed Data)] - - end - - Communities <--> |class Account| Script - Script --> |SHA256| RawDataLocal - RawDataLocal --> |Airbyte| RawDataIFT - RawDataIFT --> |dbt| ProcessedDataIFT + SDK --> SIGNAL + SDK --> |Port 8080| RPC + SDK --> |Port 8080| HTTP + SIGNAL --> |Port 8080| SOCKET + SDK --> Vol + RPC --> |coingecko_api_key| COINGECKO + RPC --> |infura_token| EVM ``` -# Setup +## Setup -## Environment Variables +To access Python funcitonality you will have to set up [Status Backend](https://github.com/status-im/status-go/). Easiest and fastest way to get it running would be with [Docker](https://www.docker.com/products/docker-desktop/). -- `POSTGRES_USERNAME` - Postgres username. -- `POSTGRES_PASSWORD` - Postgres password. -- `POSTGRES_DATABASE` - The database name in the Postgres connection. -- `POSTGRES_HOST` - The Postgres host name that will be remotely connected to. -- `POSTGRES_PORT` - The Postgres port that will be remotely connected to. -- `STATUS_DISPLAY_NAME` - The Status display name that will be used to create an account. -- `STATUS_PASSWORD` - The Status password that will be used to create an account. -- `STATUS_MNEMONIC` - The mnemonic used to recover the account. If passed a `.bkp` file will be loaded as well. Use this when you want to login to a bot account via Status App, join a community / leave community and export the `.bkp` file. -- `STATUS_INFURA_TOKEN` - [Infura token](https://www.infura.io/) is required for **token gated communities** -- `STATUS_COINGECKO_API_KEY` - [Coingecko API Key](https://www.coingecko.com/) is required for **token gated communities** - -## Docker deployement - -You can use the `docker-compose.yaml` to run the project. - -Example of `.env` file to use -``` -# Status Backend connection -STATUS_DISPLAY_NAME = "bot-status" -STATUS_PASSWORD = "ChangeThisPassword" -STATUS_MNEMONIC= "test test test test test test test test test test test test" - -# Necessary for communities that have tokens -STATUS_INFURA_TOKEN = "Token from https://www.infura.io/" -STATUS_COINGECKO_API_KEY = "Token from https://www.coingecko.com/" - -# Database config -POSTGRES_HOST=database -POSTGRES_PORT=5432 -POSTGRES_DATABASE=status-bot -POSTGRES_USERNAME=status -POSTGRES_PASSWORD=ChangeThisOneAlso +```mermaid +sequenceDiagram + actor User + participant Docker + participant Python@{"alias": "status-im/status-bot"} + participant Github@{"alias": "status-im/status-go" } + + User ->> Docker: docker-compose up + Docker ->> Github: Fetch Image + Docker ->> Docker: Build + User ->> Docker: Run container + User ->> Python: initialize module + Note over User,Python: from bot import Account
account = Account() ``` -## Python +### Docker + +Setup [`status-im/status-go`](https://github.com/status-im/status-go/) with the provided `docker-compose.yaml` file. + +``` +docker compose up -d +``` + +**Note**: To run on Windows, please make sure you clone `status-im/status-go` and change the context to the folder. If you do not want to clone the repository, make sure you have set up [WSL](https://learn.microsoft.com/en-us/windows/wsl/install) and started it. + +### Python 1. Setup environment. [Conda](https://www.anaconda.com/) example: ```bash -conda create -n status-monitoring python=3.12 +conda create -n status-sdk python=3.12 ``` **Note**: Code has been tested with **Python 3.12**. -2. Install `monitor.py` and `bot` requirements +2. Install requirements ```bash -# To run Status bot -pip install -r ./bot/requirements.txt - -# To run monitor.py pip install -r ./requirements.txt ``` - -**Note**: If you are on Windows, you will have to install `psycopg2` instead of `psycopg2-binary`. - -# Backups - -If you have already created a Status account and want to use it with it's current data, please make sure you export the `.bkp` file and put it in folder **backups** and have the following `.env` variables: - -- `STATUS_DISPLAY_NAME` -- `STATUS_PASSWORD` -- `STATUS_MNEMONIC` - -## Files - -- `monitor.py` - Status community message monitoring. It will download and upload messages in parallel. - -# Guidelines - -Things to keep in mind when building projects: - -1. **Wrong recovery phrase** can create a new account by accident. To recover the account you must correctly write the phrase. - -2. **Dynamic** `chats`, `communities` and `contacts` properties. If you log in to the account without a backup the properties will be empty **but will be populated automatically** if a message in a chat or community appears. - -3. **Display name is not the same as username**. In Status App a display name is the **curent username** of active account. This means that if you have the same profile logged in with Python, Status Desktop and Status Mobile you can have 3 different usernames **based on your current device**. If an account is logged in from more than one device, the display name will change based on the currently used one (when messaging). This feature can be handy to distinguish if a bot or an actual user is logged in to the account. **Profile pictures work in a similar way to display names** - -4. Community join requests do not work properly due to `status-go` issues. For latest updates, please monitor [`status-im/status-bot` issue #9](https://github.com/status-im/status-bot/issues/9). The fastest way to get in to a community is to manually log in with the Bot account and send a via Status App. Once the account is accepted, the bot can log in and start running. The community and chat information will appear in (2) as the messages come in. - -5. To monitor for new messages, the account **must always be logged in**. Log out of the account only if you do not want to receive messages. - -6. **Token gated chats have an impact on the entire community**. You must provide an [Infura Token](https://www.infura.io/) and [Coingecko API key](https://www.coingecko.com/) during `login`. If wallet credentials are left out, then the community will not appear in (2) properties instead of the token gated chats only. diff --git a/assets/profile.jpg b/assets/profile.jpg deleted file mode 100644 index 616f769..0000000 Binary files a/assets/profile.jpg and /dev/null differ diff --git a/assets/profile.png b/assets/profile.png new file mode 100644 index 0000000..dbe502c Binary files /dev/null and b/assets/profile.png differ diff --git a/bot/__init__.py b/bot/__init__.py index d607e0e..301bf3a 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -1,2 +1 @@ from .account import Account -from .logger import Logger diff --git a/bot/requirements.txt b/bot/requirements.txt deleted file mode 100644 index 7ab5331..0000000 --- a/bot/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ -requests -websocket-client -websockets -pandas -pillow diff --git a/config.yaml b/config.yaml deleted file mode 100644 index 35f763b..0000000 --- a/config.yaml +++ /dev/null @@ -1,23 +0,0 @@ -postgres: - schema: "status_app_monitoring" - tables: - messages: "raw_messages" - community: "raw_community_info" -# Value must be in MINUTES -sleep: 10 -files: - # Get the latest community chat dates for next run - current_state: "dates.pkl" - -bot: - # Public information for the bot - public_key: "0x041658626a9e1303b631f6d0fb1e047211d5603b977454f7d5d29fe583c3d6c1bd3d8e395d67f6c44b5bc659aae912040e9dd8164b5107368a29029cb53389d8b0" - compressed_key: "zQ3shNv1tnajHo5FvCvP662cWcbBfS5ZejB4TWaH9iAuFCZZe" - # Parameters based on Account class - params: - # localhost - to run code locally with Status Backend Dockerfile - # status-backend - to run code with docker-compose.yaml - # domain: "status-backend" - domain: "status-backend" - port: 8080 - is_secure: false diff --git a/docker-compose.yaml b/docker-compose.yaml index b30370b..806b081 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -1,10 +1,8 @@ ---- services: backend: - image: harbor.status.im/bi/status-backend:dev - # build: - # context: https://github.com/status-im/status-go.git#develop - # platform: linux/amd64 + build: + context: https://github.com/status-im/status-go.git#develop + platform: linux/amd64 container_name: status-backend ports: - 8080:8080 @@ -20,39 +18,7 @@ services: healthcheck: test: ["CMD", "curl http://0.0.0.0:8080/health"] - bot: - build: - context: . - dockerfile: Dockerfile - container_name: status-monitor - depends_on: - - backend - - database - volumes: - - ./backups:/backups - env_file: - - .env - networks: - - status-bridge - - database: - image: postgres:15 - container_name: database - env_file: - - .env - ports: - - 5432:5432 - volumes: - - postgres-data:/var/lib/postgresql - networks: - - status-bridge - healthcheck: - test: ["CMD-SHELL", "pg_isready POSTGRES_USERNAME --dbname POSTGRES_DATABASE"] - networks: status-bridge: name: status-bridge driver: bridge - -volumes: - postgres-data: diff --git a/docs/docker-compose.yaml b/docs/docker-compose.yaml deleted file mode 100644 index f29b418..0000000 --- a/docs/docker-compose.yaml +++ /dev/null @@ -1,27 +0,0 @@ -services: - backend: - # NOTE: To run on Windows, please make sure you clone `status-im/status-go` - # and change the context to the folder. If you do not want to clone the - # repository, make sure you have set up WSL (https://learn.microsoft.com/en-us/windows/wsl/install) - # and started it. - build: - context: https://github.com/status-im/status-go.git#develop - platform: linux/amd64 - container_name: status-backend - ports: - - 8080:8080 - - 8545:8545 - - 30303:30303 - entrypoint: 'status-backend' - command: '-address 0.0.0.0:8080' - volumes: - - ./backups:/root/.config/Status/backups - networks: - - status-bridge - healthcheck: - test: ["CMD", "curl http://0.0.0.0:8080/health"] - -networks: - status-bridge: - name: status-bridge - driver: bridge diff --git a/docs/overview.md b/docs/overview.md deleted file mode 100644 index 939fb42..0000000 --- a/docs/overview.md +++ /dev/null @@ -1,64 +0,0 @@ -# Status Python SDK - -![Status Python SDK header image](./images/overview-header.png) - -The initial Python Status Backend was built with testing in mind, instead of easy developer access. The objective of this repository is to make a SDK that is: - -- **light** - as less external packages when it comes to working with Status App -- **fast** - quick to get started with Status Python -- **documented** - clear explanations of what was done and **why it was done in a specific way**. - -Currently this repository is not on [PyPi](https://pypi.org/) but will be added when core functionality has been devleoped and tested. - -## How it works - -```mermaid -graph TB - subgraph backend[status-im/status-go] - subgraph Endpoints[Network: status-bridge] - RPC[RPC] - HTTP[REST] - SOCKET[Web Socket] - end - Vol[(Backup)] - end - - - subgraph bot[Python SDK] - REQUIREMENTS[requirements.txt] - SDK[class Account] - SIGNAL[class Signal] - end - - subgraph external[External Services] - COINGECKO[CoinGecko] - EVM - end - - SDK --> SIGNAL - SDK --> |Port 8080| RPC - SDK --> |Port 8080| HTTP - SIGNAL --> |Port 8080| SOCKET - SDK --> Vol - RPC --> |coingecko_api_key| COINGECKO - RPC --> |infura_token| EVM -``` - -## Setup - -To access Python funcitonality you will have to set up [Status Backend](https://github.com/status-im/status-go/). Easiest and fastest way to get it running would be with [Docker](https://www.docker.com/products/docker-desktop/). - -```mermaid -sequenceDiagram - actor User - participant Docker - participant Python@{"alias": "status-im/status-bot"} - participant Github@{"alias": "status-im/status-go" } - - User ->> Docker: docker-compose up - Docker ->> Github: Fetch Image - Docker ->> Docker: Build - User ->> Docker: Run container - User ->> Python: initialize module - Note over User,Python: from bot import Account
account = Account() -``` diff --git a/monitor.py b/monitor.py deleted file mode 100644 index 3cc105d..0000000 --- a/monitor.py +++ /dev/null @@ -1,325 +0,0 @@ -import datetime, os, pickle, yaml, time -import pandas as pd -from typing import Any -from pathlib import Path -from dotenv import load_dotenv -from hashlib import sha256 -# Manual file imports -from bot import Account, Logger -from postgres import Postgres - -def to_sha256_hash(value: str) -> str: - """ - Hash personal information before it is put in the database. - - Parameters: - - `value` - personal information - - Output: - - sha256 hashed value - """ - return sha256(value.encode()).hexdigest() - -def to_midnight(timestamp: datetime.datetime) -> datetime.datetime: - """ - Convert the given timestamp to midnight - - Parameters: - - `timestamp` - current timestap - - Output: - - `timestamp` at midnight - """ - return timestamp.replace(minute=0, second=0, hour=0, microsecond=0) - -def load_config(file_path: str) -> dict: - """ - Load the config file and the `.env` variables - - Parameter: - - `file_path` - the file path of the config yaml file. The `.env` variable must be in the same folder - - Output: - - The config variables and secret from `.env` - """ - with open(file_path, "r") as f: - config: dict = yaml.safe_load(f) - - env_file_path = os.path.join(os.path.dirname(file_path), ".env") - load_dotenv(env_file_path) - - config["env_vars"] = { - key: value - for key, value in os.environ.items() - if key.startswith(("POSTGRES_", "STATUS_")) - } - - return config - -def extract_community_channels(account: Account, community: dict, latest_dates: dict[str, pd.Timestamp]) -> pd.DataFrame: - """ - Extract the community channel messages. - - Parameters: - - `account` - logged in Status Bot account - - `community` - the current community from `account` - - `start_timestamp` - start timestamp for message fetching - - `end_timestamp` - end timestamp for message fetching - - Output: - - DataFrame with all of the community messages for the given start and end timestamps - """ - # Column name -> True if data should be hashed - bridge_key = "bridge_message" - columns = { - "id": True, - "whisper_timestamp": False, - "from": True, - "seen": False, - "chat_id": False, - "community_id": False, - "message_type": False, - "response_to": True, - "timestamp": False, - "deleted": False, - "extracted_timestamp": False, - } - - final = [] - for channel in community["channels"]: - - now = datetime.datetime.now() - start_timestamp = latest_dates.get(channel["chat_id"]) - if start_timestamp: - start_timestamp += datetime.timedelta(seconds=1) - else: - # Node will only return known / fetched messages for this channel. - # Without enabling community archives feature the node can only fetch last 30 days (from store nodes). - start_timestamp = to_midnight(now - datetime.timedelta(days=30)) - - account.logger.info(f"Starting message extraction for # {channel['name']} [{start_timestamp} - {now}]") - messages = account.get_messages(channel["chat_id"], start_timestamp, now) - messages = pd.DataFrame(messages) - if len(messages) == 0: - account.logger.info(f"No messages found") - continue - - account.logger.info(f"Extracted {len(messages)} message(s)") - messages = messages.assign( - community_id = community["id"], - extracted_timestamp = now - ) - final.append(messages) - - extracted_data = pd.concat(final, ignore_index=True) if final else pd.DataFrame() - if len(extracted_data) == 0: - return extracted_data - - existing_columns = extracted_data.columns.to_list() - for column, should_hash in columns.items(): - if column not in existing_columns: - loc = len(extracted_data.columns.to_list()) - extracted_data.insert(loc, column, None) - continue - - if should_hash: - extracted_data[column] = extracted_data[column].astype(str).apply(to_sha256_hash) - - if bridge_key in extracted_data.columns: - extracted_data["source"] = extracted_data[bridge_key].apply(lambda value: value["bridgeName"] if not pd.isna(value) else "status") - else: - extracted_data["source"] = "status" - - extracted_data = extracted_data[list(columns.keys()) + ["source"]].assign( - deleted = extracted_data["deleted"].fillna(False), - seen = extracted_data["seen"].fillna(False) - ) - account.logger.info(f"Sensitive data has been hashed") - - return extracted_data - -def save_file(file_path: str, data: Any): - """ - Save data to a pickle file. Creates directories if they don't exist. - - Parameters: - - `file_path` - Full pikle path - - `data` - Python object to be saved - """ - folder = os.path.dirname(file_path) - if len(folder) > 0: - os.makedirs(folder, exist_ok=True) - - if isinstance(data, pd.DataFrame): - data.to_csv(file_path, index=False) - return - - with open(file_path, "wb") as f: - pickle.dump(data, f) - -def create_bot(config: dict) -> Account: - """ - Initialized a logged in bot account that will monitor the communities. - - Parameters: - - `config` - the `load_config` configuration - - Output: - - Logged in Bot account - """ - params = config.get("bot", {}).get("params", {}) - account = Account(**params) - available_accounts = [acc["display_name"] for acc in account.available_accounts] - - prefix = "STATUS_" - params = { - key.replace(prefix, "").lower(): value - for key, value in config["env_vars"].items() - if key.startswith(prefix) - } - if params["display_name"] in available_accounts: - params.pop("mnemonic") - - account.login(**params) - if account.info["compressed_key"] != config["bot"]["compressed_key"]: - raise Exception("Target compressed key and logged in compressed key are different...") - else: - account.logger.info("[SUCCESS] Logged in with correct account") - - balance = account["GBP"] - query = (balance["symbol"] == "SNT") & (balance["fiat_value"] > 0) & (balance["chain_id"] == 1) - if query.sum() != 1: - raise Exception("There were issues with Infura Token and Coingecko initialization...") - else: - account.logger.info("[SUCCESS] Wallet balance is available") - - account.profile_picture = os.path.join(os.path.dirname(__file__), "assets", "profile.jpg") - account.logger.info(f"Account Information:\nCompressed Key: {account.info['compressed_key']}\nPublic Key: {account.info['public_key']}\nURL: {account.info['url']}") - return account - -def download(account: Account, folder: str, config: dict): - """ - Download Status App messages / info from communities and store them in pickle files. - - Parameters: - - `folder` - the folder where the files will be created. Sub folders are automatically created - - `config` - the `load_config` configuration - """ - file_path = os.path.join(os.path.dirname(__file__), config["files"]["current_state"]) - latest_dates: dict[str, pd.Timestamp] = pd.read_pickle(file_path) if os.path.exists(file_path) else {} - - get_file_name = lambda: str(to_midnight(datetime.datetime.now()).timestamp()).replace(".", "") - communities = account.communities - if not communities: - account.logger.warning("No communities found...") - - for community in communities: - - if not community["is_member"]: - continue - - community_folder_name = community["name"].replace(" ", "-") - messages_folder = os.path.join(folder, "messages", community_folder_name) - community_info_folder = os.path.join(folder, "community", community_folder_name) - - account.logger.info(f"Extracting data for {community['name']}") - community["extracted_timestamp"] = datetime.datetime.now() - - file_path = os.path.join(community_info_folder, get_file_name() + ".pkl") - if not os.path.exists(file_path): - save_file(file_path, community) - account.logger.info(f"Created {file_path}") - - file_path = os.path.join(messages_folder, get_file_name() + ".csv") - if not os.path.exists(file_path): - messages = extract_community_channels(account, community, latest_dates) - if len(messages) > 0: - save_file(file_path, messages) - account.logger.info(f"Created {file_path}") - -def store(folder: str, config: dict, logger: Logger): - """ - Upload Status App `download` file to Postgres. - NOTE: The Postgres schema must already exist - - Parameters: - - `folder` - the folder where the files will be created. Sub folders are automatically created - - `config` - the `load_config` configuration - """ - path = Path(folder) - table_name_mapping: dict[str, str] = config["postgres"]["tables"] - table_schema = config["postgres"]["schema"] - - upload: dict[str, list] = {} - - file_path = os.path.join(os.path.dirname(__file__), config["files"]["current_state"]) - latest_dates: dict[str, pd.Timestamp] = pd.read_pickle(file_path) if os.path.exists(file_path) else {} - - completed = [] - - files = list(path.rglob("*.pkl")) + list(path.rglob("*.csv")) - logger.info(f"There are {len(files)} file(s) to upload") - for file_path in files: - - table_name = table_name_mapping.get(file_path.parent.parent.name) - if not table_name: - continue - - file_name = str(file_path.name) - data = pd.read_pickle(file_path) if file_name.endswith(".pkl") else pd.read_csv(file_path) - if isinstance(data, dict): - data = pd.DataFrame([data]) - - for column in data.columns: - if "timestamp" not in column: - continue - data[column] = pd.to_datetime(data[column]) - - if table_name not in upload: - upload[table_name] = [] - - if "timestamp" in data.columns: - latest_dates.update(data.groupby("chat_id")["timestamp"].max().to_dict()) - - upload[table_name].append(data) - completed.append(str(file_path)) - - save_file(config["files"]["current_state"], latest_dates) - logger.info(f"Updated {config['files']['current_state']}") - - prefix = "POSTGRES_" - params = { - key.replace(prefix, "").lower(): value - for key, value in config["env_vars"].items() - if key.startswith(prefix) - } - connector = Postgres(**params) - for table_name, data in upload.items(): - if len(data) == 0: - continue - - df = pd.concat(data, ignore_index=True).assign(batch_timestamp = datetime.datetime.now()) - json_columns = [ - column - for column in df.columns - if len(df[column].dropna()) > 0 and isinstance(df[column].dropna().reset_index(drop=True).iloc[0], (dict, list)) - ] - connector.insert(df, table_name, table_schema, json_columns) - logger.info(f"Uploaded {len(df)} record(s) to {table_schema}.{table_name}") - - for file_path in completed: - os.remove(file_path) - logger.info(f"Deleted {file_path}") - -if __name__ == "__main__": - folder = os.path.dirname(__file__) - config = load_config(os.path.join(folder, "config.yaml")) - upload_folder = os.path.join(os.path.dirname(__file__), "uploads") - logger = Logger() - account = create_bot(config) - - while True: - download(account, upload_folder, config) - store(upload_folder, config, logger) - logger.info(f"Sleeping for {config['sleep']} minute(s)") - time.sleep(config["sleep"] * 60) diff --git a/postgres.py b/postgres.py deleted file mode 100644 index 6a4512e..0000000 --- a/postgres.py +++ /dev/null @@ -1,149 +0,0 @@ -""" -Minimum code to upload data taken from: -https://github.com/status-im/ift-data-py/blob/master/ift_data/clients/postgres.py -""" - -import psycopg2 -import pandas as pd -from typing import Optional, Union -from sqlalchemy import create_engine -from sqlalchemy.dialects.postgresql import JSONB - -class Postgres: - - def __init__(self, username: str, password: str, port: Union[int, str], database: str, host: str): - - if isinstance(port, str): - port = int(port) - - self.__params = { - "host": host, - "user": username, - "password": password, - "port": port, - "database": database - } - - self.__url = f"postgresql://{username}:{password}@{host}:{port}/{database}" - self.__conn: psycopg2.extensions.connection = psycopg2.connect(**self.__params) - self.__cursor: psycopg2.extensions.cursor = self.__conn.cursor() - - def insert(self, data: pd.DataFrame, table_name: str, schema: str, json_columns: Optional[list] = None): - """ - Insert the DataFrame in the specified schema > table. - If the schema / table name does not exist, it will be created. - - Parameters: - - `data` - the data to be inserted in Postgres - - `table_name` - the name of the table - - `schema` - the name of the schema - - `json_columns` - when creating the table, `dict` columns will be turned into JSON objects in Postgres - """ - self.execute(f"CREATE SCHEMA IF NOT EXISTS {schema}") - engine = create_engine(self.__url) - - data.columns = [column.lower() for column in data.columns] - - params = { - "name": table_name, - "con": engine, - "schema": schema, - "if_exists": "append", - "index": False - } - if json_columns: - params["dtype"] = { - json_column: JSONB - for json_column in json_columns - } - - # Add new columns as they come - existing_columns = self.get_columns(schema, table_name) - - if existing_columns: - for column in data.columns: - if column in existing_columns: - continue - # NOTE: New values will have to be transformed - self.execute(f"ALTER TABLE {schema}.{table_name} ADD COLUMN {column} TEXT") - - data.to_sql(**params) - - def execute(self, query: str): - """ - Execute queries such as INSERT, UPDATE, DELETE etc. - - Parameters: - - `query` - the PostgreSQL query - """ - self.__execute(query) - self.__conn.commit() - - def to_pandas(self, query: str, batch_size: int = 50_000, uppercase: bool = True) -> pd.DataFrame: - """ - Create a DataFrame from the given query - - Parameters: - - `query` - the PostgreSQL query - - `batch_size` - how many rows will be fetched at once - - `uppercase` - if `True` then the columns will be uppercase. If `False` the columns will be lowercase - Output: - - DataFrame for the executed query - """ - self.__execute(query) - columns = [column.name.upper() if uppercase else column.name.lower() for column in self.__cursor.description] - chunks = [] - - while True: - rows = self.__cursor.fetchmany(batch_size) - if not rows: - break - chunks.append(pd.DataFrame(rows, columns=columns)) - - return pd.concat(chunks, ignore_index=True) if chunks else pd.DataFrame(columns=columns) - - def close(self): - self.__cursor.close() - self.__conn.close() - - def __del__(self): - self.close() - - def __execute(self, query: str): - - failed = False - is_closed = bool(self.__conn.closed) - - if is_closed: - self.__conn: psycopg2.extensions.connection = psycopg2.connect(**self.__params) - self.__cursor: psycopg2.extensions.cursor = self.__conn.cursor() - - try: - self.__cursor.execute(query) - except psycopg2.errors.InFailedSqlTransaction: - self.__conn.rollback() - failed = True - - if failed: - self.__cursor.execute(query) - - - def get_columns(self, schema: str, table_name: str) -> list[str]: - """ - Get the column names in the correct order for the given table. - - Parameters: - - `table_name` - the name of the table - - `schema` - the name of the schema - - Output: - - the table's columns in the correct order - """ - query = f""" - SELECT column_name - FROM information_schema.columns - WHERE table_name = '{table_name}' - AND table_schema = '{schema}' - ORDER BY ordinal_position ASC - """ - return self.to_pandas(query)["COLUMN_NAME"].to_list() diff --git a/requirements.txt b/requirements.txt index 3307083..7ab5331 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ +requests +websocket-client +websockets pandas -pyyaml -python-dotenv -psycopg2-binary -sqlalchemy +pillow