import logging
from pgvector.psycopg2 import register_vector
import psycopg2


class Database:
    def __init__(
        self,
        database,
        user="postgres",
        password="postgres",
        host="localhost",
        port="5433",
    ) -> None:
        logging.info("Connecting to database")
        self.conn = psycopg2.connect(
            database="rag",
            user="postgres",
            password="postgres",
            host="localhost",
            port="5433",
        )
        register_vector(self.conn)
        self.cur = self.conn.cursor()
        self.cur.execute("SELECT version();")
        logging.info(" DB Version: %s", self.cur.fetchone()[0])
        logging.info(" psycopg2 Version: %s", psycopg2.__version__)

    def close(self):
        logging.info("Closing connection to database")
        self.cur.close()
        self.conn.close()