shakecities/sites/db.py

import mysql.connector
import os
import dotenv
import json

dotenv.load_dotenv()

# Database connection
dbargs = {
    'host':os.getenv('DB_HOST'),
    'user':os.getenv('DB_USER'),
    'password':os.getenv('DB_PASSWORD'),
    'database':os.getenv('DB_NAME')
}

def get_website_data(domain):
    connection = mysql.connector.connect(**dbargs)
    cursor = connection.cursor()
    cursor.execute("""
        SELECT * FROM site WHERE domain = %s
    """, (domain,))
    data = cursor.fetchall()
    cursor.close()
    connection.close()
    
    if data == []:
        return ""
    
    parsed = data[0][2]
    parsed = json.loads(parsed)
    parsed = parsed['data']
    # Decoding
    parsed = parsed.encode('utf-8').decode('unicode-escape')

    return parsed

def get_website_data_raw(domain):
    connection = mysql.connector.connect(**dbargs)
    cursor = connection.cursor()
    cursor.execute("""
        SELECT * FROM site WHERE domain = %s
    """, (domain,))
    data = cursor.fetchall()
    cursor.close()
    connection.close()
    
    if data == []:
        # Create new entry
        connection = mysql.connector.connect(**dbargs)
        cursor = connection.cursor()
        data = {
            "data": ""
        }
        insert_query = "INSERT INTO site (data,domain) VALUES (%s,%s)"
        cursor.execute(insert_query, (json.dumps(data), domain))
        connection.commit()
        cursor.close()
        connection.close()
        return ""
    
    parsed = data[0][2]
    parsed = json.loads(parsed)

    return parsed

def get_website_wallet(domain,token):
    connection = mysql.connector.connect(**dbargs)
    cursor = connection.cursor()
    cursor.execute("""
        SELECT * FROM site WHERE domain = %s
    """, (domain,))
    data = cursor.fetchall()
    cursor.close()
    connection.close()
    
    if data == []:
        return ""
    
    parsed = data[0][2]
    parsed = json.loads(parsed)
    if token in parsed:
        parsed = parsed[token]
        return parsed
    
    return ""