@Wolfengo
Я ещё зелёный, сильно не душите

Как правильно сделать декоратор для открытия подключения к базе данных (на основе моего кода с комментариями)?

Как правильно сделать декоратор для открытия подключения к базе данных (на основе моего кода с комментариями)?

import datetime
import psycopg2
from psycopg2 import errors


class ConnectDecorator:
    def __init__(self, dbname, user, password, host, port):
        self.dbname = dbname
        self.user = user
        self.password = password
        self.host = host
        self.port = port
        self.iteration = 0

    # При вызове класса будет вызываться декоратор для открытия
    def __call__(self, cls):
        for name, value in vars(cls).items():
            if callable(value):
                setattr(cls, name, self.decorate_method(value))
        return cls

    # Декоратор
    def decorate_method(self, func):
        def wrapper(*args, **kwargs):
            conn = None
            # Если подключение не получилось, то уходит создавать новую таблицу и добавляет итерацию +1
            try:
                conn = psycopg2.connect(dbname=self.dbname, user=self.user, password=self.password, host=self.host)
                # Проверяет, если до этого база данных создавалась, то создаёт для неё таблицу. 
                # Если итерация 0, то идёт дальше
                if self.iteration == 1:
                    cursor = conn.cursor()
                    cursor.execute(f'CREATE TABLE IF NOT EXISTS public.users '
                                   f"(id integer NOT NULL, user_id integer NOT NULL, PRIMARY KEY (id));"
                                   f'ALTER TABLE IF EXISTS public.users OWNER to postgres;')
                    print(f"[INFO] {datetime.datetime.now()} Таблица пользователей успешно создана")
                    conn.commit()
            except errors.OperationalError:
                if self.iteration < 1:
                    try:
                        # Создание базы данных
                        self.create_database_if_not_exists()
                        self.iteration += 1
                        # Создаём рекурсию, чтобы заново пробежаться по функции
                        self.decorate_method(func)
                    except Exception as e:
                        print(f"Ошибка создания новой базы данных: {e}")
            try:
                return func(conn, *args, **kwargs)
            except Exception as e:
                print(e)
            finally:
                if conn is not None:
                    conn.close()

        return wrapper

    def create_database_if_not_exists(self):
        # Подключаемся к базе postgres, чтобы узнать, существует ли база
        conn = psycopg2.connect(database='postgres', user=self.user, password=self.password, host=self.host,
                                port=self.port)
        conn.autocommit = True
        cursor = conn.cursor()

        # Проверка наличия базы данных
        cursor.execute("SELECT 1 FROM pg_database WHERE datname = %s", (self.dbname,))
        exists = cursor.fetchone()

        if not exists:
            try:
                # Создание новой базы данных
                with conn.cursor() as cursor:
                    cursor.execute(f"CREATE DATABASE {self.dbname}")
                    print(f"База данных '{self.dbname}' создана.")
            except Exception as e:
                print(f'Ошибка {e}')
            finally:
                if conn:
                    conn.close()
        else:
            print(f"База данных '{self.dbname}' уже существует.")

        cursor.close()
        conn.close()


# Декоратор для открытия подключения
@ConnectDecorator(dbname='mydatabase', user='postgres', password='123', host='localhost', port='5432')
class PostgreSQLDatabase:
    def user_exists(self):
        """Проверяем наличие в базе данных"""
        try:
            cursor = self.conn.cursor()
            cursor.execute("SELECT * FROM users")
            print(f"[INFO] {datetime.datetime.now()} Данные успешно получены")
            return bool(len(cursor.fetchall()))
        except Exception as e:
            print(f"[ERROR] Ошибка работы базы данных: {e}")


db = PostgreSQLDatabase()
db.user_exists()
  • Вопрос задан
  • 257 просмотров
Решения вопроса 1
@Wolfengo Автор вопроса
Я ещё зелёный, сильно не душите
По совету из комментариев, реализовал код вот так:

import psycopg2.pool
from psycopg2 import errors

dbname = 'mydatabase'
user = 'postgres'
password = '123'
host = 'localhost'
port = '5432'
connection_string = f"dbname={dbname} user={user} password={password} host={host} port={port}"


class DatabaseConnection:
    def __init__(self, connection_string):
        self._connection_string = connection_string
        self._connection_pool = psycopg2.pool.SimpleConnectionPool(minconn=1, maxconn=10, dsn=connection_string)
        self._connection = None

    def __enter__(self):
        self._connection = self._connection_pool.getconn()
        return self._connection.cursor()

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self._connection:
            self._connection.commit()
            self._connection_pool.putconn(self._connection)


class PostgreSQLDatabase:
    def __init__(self, connection_string):
        self._connection_string = connection_string

    def __enter__(self):
        self._connection = DatabaseConnection(self._connection_string)
        return self._connection.__enter__()

    def __exit__(self, exc_type, exc_val, exc_tb):
        return self._connection.__exit__(exc_type, exc_val, exc_tb)


class Admin_connected:
    def __init__(self):
        self.database = PostgreSQLDatabase(connection_string)

    def create_database_if_not_exists(self):
        # Подключаемся к базе postgres, чтобы узнать, существует ли база
        conn = psycopg2.connect(database='postgres', user=user, password=password, host=host,
                                port=port)
        conn.autocommit = True
        cursor = conn.cursor()

        # Проверка наличия базы данных
        cursor.execute("SELECT 1 FROM pg_database WHERE datname = %s", (dbname,))
        exists = cursor.fetchone()

        if not exists:
            try:
                # Создание новой базы данных
                with conn.cursor() as cursor:
                    cursor.execute(f"CREATE DATABASE {dbname}")
                    print(f"База данных '{dbname}' создана.")
                    self.create_user_table()
            except Exception as e:
                print(f'Ошибка {e}')
            finally:
                if conn:
                    conn.close()
        else:
            print(f"База данных '{dbname}' уже существует.")

    def create_user_table(self):
        try:
            with self.database as cursor:
                cursor.execute(f'CREATE TABLE IF NOT EXISTS public.users '
                               f"(id integer NOT NULL, user_id integer NOT NULL, PRIMARY KEY (id));"
                               f'ALTER TABLE IF EXISTS public.users OWNER to postgres;')
                print(f"[INFO] Таблица пользователей успешно создана")
        except Exception as e:
            print(f'Ошибка при создании новой таблицы: {e}')

    def exists(self):
        """Проверяем на наличие в базе данных"""
        try:
            with self.database as cursor:
                cursor.execute("SELECT * FROM users")
                print("[INFO] Данные успешно получены")
                return cursor.fetchall()
        except errors.OperationalError:
            self.create_database_if_not_exists()
            return self.exists()


db = Admin_connected()
print(db.exists())
Ответ написан
Комментировать
Пригласить эксперта
Ваш ответ на вопрос

Войдите, чтобы написать ответ

Похожие вопросы