MenuVision-AI / modules /database.py
Wall06's picture
Create modules/database.py
3bc2a85 verified
raw
history blame
1.01 kB
import sqlite3
import pandas as pd
DB_NAME = "menuvision.db"
def init_db():
conn = sqlite3.connect(DB_NAME)
c = conn.cursor()
# Create Tables
c.execute('''CREATE TABLE IF NOT EXISTS restaurants
(id INTEGER PRIMARY KEY, name TEXT, location TEXT, owner TEXT)''')
conn.commit()
conn.close()
def add_restaurant(name, location, owner):
try:
conn = sqlite3.connect(DB_NAME)
c = conn.cursor()
c.execute("INSERT INTO restaurants (name, location, owner) VALUES (?, ?, ?)", (name, location, owner))
conn.commit()
conn.close()
return f"Success: Restaurant '{name}' added!"
except Exception as e:
return f"Error: {e}"
def search_restaurants(query):
conn = sqlite3.connect(DB_NAME)
# Simple SQL query to find matches
sql = f"SELECT name, location, owner FROM restaurants WHERE name LIKE '%{query}%' OR location LIKE '%{query}%'"
df = pd.read_sql_query(sql, conn)
conn.close()
return df