status-im/status-python-sdk
1
0
Fork 0
mirror of https://github.com/status-im/status-python-sdk.git synced 2026-08-30 21:51:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

bot: Remove monitor.py

Browse Source
This commit is contained in:
Nick Ninov
2026-05-24 17:35:20 +03:00
parent e3e338ea21
commit f97f6b66b6
15 changed files with 66 additions and 842 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.dockerignore
-14
View File
@@ -1,14 +0,0 @@
# Git
.git
.gitignore
# Python
__pycache__/
*.pyc
# Status Bot
uploads/
bot/docs
config.yaml
.env
Dockerfile
-13
View File
@@ -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 []
Jenkinsfile
Vendored
-71
View File
@@ -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() }
}
}
README.md
+59 -109
View File
@@ -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<br>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.
assets/profile.jpg
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 14 KiB

assets/profile.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 177 KiB

bot/__init__.py
-1
View File
@@ -1,2 +1 @@
from .account import Account
from .logger import Logger
bot/requirements.txt
-5
View File
@@ -1,5 +0,0 @@
requests
websocket-client
websockets
pandas
pillow
config.yaml
-23
View File
@@ -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
docker-compose.yaml
+3 -37
View File
@@ -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:
docs/docker-compose.yaml
-27
View File
@@ -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
docs/overview.md
-64
View File
@@ -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<br>account = Account()
```
monitor.py
-325
View File
@@ -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)
postgres.py
-149
View File
@@ -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()
requirements.txt
+4 -4
View File
@@ -1,5 +1,5 @@
requests
websocket-client
websockets
pandas
pyyaml
python-dotenv
psycopg2-binary
sqlalchemy
pillow