Practice/src/db.py

import psycopg2
from psycopg2.extras import DictCursor
import os
from dotenv import load_dotenv

# Список организаций, сотрудников и корреспонденций
organizations = []
employees = []
correspondencies = []

# Подключение к базе данных
conn = None

# Инициализация подключения к базе данных
def init_connection(username, password):
   global conn
   # Загрузка переменных окружения из .env файла
   load_dotenv()

   # Подключение к базе данных
   conn = psycopg2.connect(
      dbname=os.getenv("DB_NAME"),
      user=username,
      password=password,
      host=os.getenv("DB_HOST"),
      port=os.getenv("DB_PORT")
   )

   # Изменяем текущую схему
   with conn.cursor() as cur:
      cur.execute(f"SET search_path TO {os.getenv('DB_SCHEMA')};")


# Получение всех организаций
def get_all_organizations():
   global organizations
   with conn.cursor(cursor_factory=DictCursor) as cur:
      cur.execute("SELECT id, name, legal_form FROM organization;")
      organizations = cur.fetchall()

# Получение всех сотрудников
def get_all_employees():
   global employees
   with conn.cursor(cursor_factory=DictCursor) as cur:
      cur.execute("SELECT id, name, lastname, patronymic, organization_id FROM employee;")
      employees = cur.fetchall()

# Получение всех корреспонденций
def get_all_correspondencies():
   global correspondencies
   with conn.cursor(cursor_factory=DictCursor) as cur:
      cur.execute("SELECT id, document_type, document_number, receipt_date, execution_date, employee_id FROM correspondence;")
      correspondencies = cur.fetchall()

# Добавление новой корреспонденции
def add_correspondence(document_type, document_number, receipt_date, execution_date, employee_id):
   with conn.cursor() as cur:
      cur.execute(
         "INSERT INTO correspondence (document_type, document_number, receipt_date, execution_date, employee_id) VALUES (%s, %s, %s, %s, %s) RETURNING id;",
         (document_type, document_number, receipt_date, execution_date, employee_id)
      )
      new_id = cur.fetchone()[0]
      conn.commit()
      return new_id
   
# Редактирование существующей корреспонденции
def edit_correspondence(correspondence_id, document_type, document_number, receipt_date, execution_date, employee_id):
   with conn.cursor() as cur:
      cur.execute(
         "UPDATE correspondence SET document_type = %s, document_number = %s, receipt_date = %s, execution_date = %s, employee_id = %s WHERE id = %s;",
         (document_type, document_number, receipt_date, execution_date, employee_id, correspondence_id)
      )
      conn.commit()

# Удаление корреспонденции
def delete_correspondence(correspondence_id):
   with conn.cursor() as cur:
      cur.execute("DELETE FROM correspondence WHERE id = %s;", (correspondence_id,))
      conn.commit()

# Подсчитать количество корреспонденций для выбранной
def count_correspondencies_by_organization(organization_id):
   with conn.cursor() as cur:
      cur.execute(
         "SELECT COUNT(*) FROM correspondence c JOIN employee e ON c.employee_id = e.id WHERE e.organization_id = %s;",
         (organization_id,)
      )
      count = cur.fetchone()[0]
      return count