Spaces:
Sleeping
Sleeping
File size: 3,577 Bytes
9dfeb39 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | import os
import psycopg2
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
from dotenv import load_dotenv
load_dotenv()
class DatabaseSetup:
def __init__(self):
self.host = os.getenv('DB_HOST', 'localhost')
self.port = os.getenv('DB_PORT', '5432')
self.db_name = os.getenv('DB_NAME', 'cocktails_db')
self.user = os.getenv('DB_USER', 'postgres')
self.password = os.getenv('DB_PASSWORD', 'your_password')
def create_database(self):
"""Create the database if it doesn't exist"""
try:
# Connect to default postgres database
conn = psycopg2.connect(
host=self.host,
port=self.port,
user=self.user,
password=self.password,
database='postgres'
)
conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
cursor = conn.cursor()
# Check if database exists
cursor.execute(f"SELECT 1 FROM pg_catalog.pg_database WHERE datname = '{self.db_name}'")
exists = cursor.fetchone()
if not exists:
cursor.execute(f'CREATE DATABASE {self.db_name}')
print(f"Database '{self.db_name}' created successfully")
else:
print(f"Database '{self.db_name}' already exists")
cursor.close()
conn.close()
except Exception as e:
print(f"Error creating database: {e}")
def setup_pgvector(self):
"""Setup pgvector extension and create tables"""
try:
conn = psycopg2.connect(
host=self.host,
port=self.port,
user=self.user,
password=self.password,
database=self.db_name
)
cursor = conn.cursor()
# Enable pgvector extension
cursor.execute("CREATE EXTENSION IF NOT EXISTS vector")
# Create cocktails table with vector embeddings
cursor.execute("""
CREATE TABLE IF NOT EXISTS cocktails (
id SERIAL PRIMARY KEY,
name VARCHAR(255) NOT NULL,
ingredients TEXT NOT NULL,
recipe TEXT,
glass VARCHAR(100),
category VARCHAR(100),
iba VARCHAR(100),
alcoholic VARCHAR(50),
embedding vector(384)
)
""")
# Create index for vector similarity search
cursor.execute("""
CREATE INDEX IF NOT EXISTS cocktails_embedding_idx
ON cocktails USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 100)
""")
conn.commit()
cursor.close()
conn.close()
print("Database tables and pgvector extension set up successfully")
except Exception as e:
print(f"Error setting up pgvector: {e}")
def get_connection(self):
"""Get database connection"""
return psycopg2.connect(
host=self.host,
port=self.port,
user=self.user,
password=self.password,
database=self.db_name
)
if __name__ == "__main__":
db_setup = DatabaseSetup()
db_setup.create_database()
db_setup.setup_pgvector()
|