database.py

import pymysql
import streamlit as st
import logging
import json
import pandas as pd
import re
import os
from langchain_community.utilities.sql_database import SQLDatabase


db_params = {"host": os.getenv("DB_HOST"),
                "user": os.getenv("DB_USER"),
                "password": os.getenv("DB_PASSWORD"),
                "port": 3306,
                "database":os.getenv("DB")
                }



def initialize_database():
    try:
        # Database Connection
        db = pymysql.connect(**db_params)
        st.success("Database connection successful!")
        return db
    except Exception as e:
        st.error(f"Database connection failed: {e}")
        return None


def get_db():
    try:
        db = SQLDatabase.from_uri(
            f"mysql+pymysql://{db_params['user']}:{db_params['password']}@{db_params['host']}/{db_params['database']}",
            include_tables=['term_details_modified', 'veda_content_details', 'veda_content_modified']
        )
        #st.success("Database connection successful!")
        return db
    except Exception as e:
        st.error(f"Database connection failed: {e}")
        return None


def execute_query(query):
    db = initialize_database()
    cursor = db.cursor()
    try:
        cursor.execute(query)
        description = cursor.description
        result = cursor.fetchall()  # Fetch all rows from the result set
        db.commit()
        return description, result
    except Exception as e:
        print("Error executing query:", e)
        db.rollback()
        return None  # Return None if an error occurs
    finally:
        db.close()


def execute_sql_query(query, parameters=None):
    # Establish database connection and execute SQL query
    db = initialize_database()
    cursor = db.cursor(pymysql.cursors.DictCursor)  # Use dictionary cursor to retrieve data as dictionaries
    try:
        if parameters:
            cursor.execute(query, parameters)
        else:
            cursor.execute(query)
        results = cursor.fetchall()
        return results
    except Exception as e:
        logging.error(f"Error executing SQL query: {e}")
        return None
    finally:
        db.close()


def get_details_mantra_json(query):
    description, data = execute_query(query)
    df = pd.DataFrame(data)
    df.columns = [x[0] for x in description]
    mantra_json = df['mantra_json'].values[0]
    cleaned_data = re.sub('<[^<]+?>', '', mantra_json)
    return json.loads(cleaned_data)