db.py

import json
import psycopg2
import uuid
import os

def init_db():
    con = psycopg2.connect(
        dbname=os.getenv('POSTGRES_DB'),
        user=os.getenv('POSTGRES_USER'),
        password=os.getenv('POSTGRES_PASSWORD'),
        host=os.getenv('POSTGRES_HOST'),
        port=os.getenv('POSTGRES_PORT')
    )
    cur = con.cursor()

    cur.execute('''CREATE TABLE IF NOT EXISTS MarketData
                (id UUID PRIMARY KEY, symbol TEXT, interval TEXT, close REAL, high REAL, low REAL, timestamp BIGINT)''')
    con.commit()
    con.close()

def log_db(message):
    parsed = json.loads(message)

    k = parsed["k"]

    id = uuid.uuid4()  # PostgreSQL supports UUID natively
    symbol = k["s"]
    interval = k["i"]
    close = float(k["c"])
    high = float(k["h"])
    low = float(k["l"])
    time = int(k["t"])

    print("log:")
    print("symbol: " + symbol)
    print("interval: " + interval)
    print("close: " + str(close))
    print("high: " + str(high))
    print("low: " + str(low))
    print("time: " + str(time))

    row = (str(id), symbol, interval, close, high, low, time)
    con = psycopg2.connect(
        dbname=os.getenv('POSTGRES_DB'),
        user=os.getenv('POSTGRES_USER'),
        password=os.getenv('POSTGRES_PASSWORD'),
        host=os.getenv('POSTGRES_HOST'),
        port=os.getenv('POSTGRES_PORT')
    )
    cur = con.cursor()
    cmd = "INSERT INTO MarketData (id, symbol, interval, close, high, low, timestamp) VALUES (%s, %s, %s, %s, %s, %s, %s)"
    cur.execute(cmd, row)
    
    con.commit()
    con.close()


def get_recent_market_data():
    con = psycopg2.connect(
        dbname=os.getenv('POSTGRES_DB'),
        user=os.getenv('POSTGRES_USER'),
        password=os.getenv('POSTGRES_PASSWORD'),
        host=os.getenv('POSTGRES_HOST'),
        port=os.getenv('POSTGRES_PORT')
    )
    cur = con.cursor()
    cur.execute("SELECT * FROM MarketData ORDER BY timestamp ASC")
    rows = cur.fetchall()
    con.close()
    return rows

def get_recent_market_data_length():
    try:
        con = psycopg2.connect(
            dbname=os.getenv('POSTGRES_DB'),
            user=os.getenv('POSTGRES_USER'),
            password=os.getenv('POSTGRES_PASSWORD'),
            host=os.getenv('POSTGRES_HOST'),
            port=os.getenv('POSTGRES_PORT')
        )
        cur = con.cursor()
        cur.execute("SELECT COUNT(*) FROM MarketData")  # Count the number of rows
        row_count = cur.fetchone()[0]  # Fetch the count
        con.close()
        return row_count
    except psycopg2.Error as e:
        print(f"Error fetching data: {e}")
        return 0  # Return 0 in case of an error

# def get_recent_market_data(period):
#     con = psycopg2.connect(
#         dbname=os.getenv('POSTGRES_DB'),
#         user=os.getenv('POSTGRES_USER'),
#         password=os.getenv('POSTGRES_PASSWORD'),
#         host=os.getenv('POSTGRES_HOST'),
#         port=os.getenv('POSTGRES_PORT')
#     )
#     cur = con.cursor()
#     cur.execute(f"SELECT * FROM MarketData ORDER BY timestamp DESC LIMIT {period}")
#     rows = cur.fetchall()
#     con.close()
#     return rows