33 lines
880 B
Python
33 lines
880 B
Python
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()
|