Files
WFL/db.py

37 lines
897 B
Python

import psycopg2
from psycopg2.extras import RealDictCursor
from dotenv import load_dotenv
import os
load_dotenv()
def get_connection():
# return psycopg2.connect(
# host="localhost",
# port=5432,
# dbname="wfl_db",
# user="postgres",
# password="159753"
# )
return psycopg2.connect(
host=os.getenv("DB_HOST"),
port=os.getenv("DB_PORT"),
dbname=os.getenv("DB_NAME"),
user=os.getenv("DB_USER"),
password=os.getenv("DB_PASSWORD")
)
def test_connection():
conn = get_connection()
try:
with conn.cursor(cursor_factory=RealDictCursor) as cur:
cur.execute("SELECT current_database() AS db_name;")
row = cur.fetchone()
print(f"Connected to database: {row['db_name']}")
finally:
conn.close()
if __name__ == "__main__":
test_connection()