diff --git a/site/app.py b/site/app.py index 062f8d3..70c3e9e 100644 --- a/site/app.py +++ b/site/app.py @@ -9,6 +9,7 @@ load_dotenv() class ContractsApp: def __init__(self): self.app = Flask(__name__) + db.ensure_db() self.add_routes() def add_routes(self): diff --git a/site/db.py b/site/db.py index 2a22546..5618fac 100644 --- a/site/db.py +++ b/site/db.py @@ -3,22 +3,47 @@ import psycopg2 from psycopg2 import sql +def _pg_connect(dbname): + """Сырое подключение к указанной базе.""" + return psycopg2.connect( + host=os.getenv("DB_HOST"), + port=os.getenv("DB_PORT", 5432), + dbname=dbname, + user=os.getenv("DB_USER"), + password=os.getenv("DB_PASS"), + sslmode=os.getenv("DB_SSLMODE", "disable"), + ) + + def connect(): - """Подключение к БД. Возвращает (connection, None) или (None, error).""" + """Подключение к целевой БД. Возвращает (connection, None) или (None, error).""" try: - conn = psycopg2.connect( - host=os.getenv("DB_HOST"), - port=os.getenv("DB_PORT", 5432), - dbname=os.getenv("DB_NAME"), - user=os.getenv("DB_USER"), - password=os.getenv("DB_PASS"), - sslmode=os.getenv("DB_SSLMODE", "disable"), - ) - return conn, None + return _pg_connect(os.getenv("DB_NAME")), None except Exception as e: return None, str(e) +def ensure_db(): + """Создать базу, если не существует.""" + db_name = os.getenv("DB_NAME", "contracts") + try: + conn = _pg_connect("postgres") + conn.autocommit = True + cur = conn.cursor() + cur.execute("SELECT 1 FROM pg_database WHERE datname = %s", (db_name,)) + if cur.fetchone(): + print(f"DB '{db_name}' exists") + else: + cur.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(db_name))) + print(f"DB '{db_name}' created") + cur.close() + conn.close() + return True + except Exception as e: + print(f"ensure_db error: {e}") + return False + + def query(sql_text, params=None): """Выполнить запрос, вернуть (result, error).""" conn, err = connect()