Переделал метод запроса в бд

page_analyzer/db.py

@@ -24,7 +24,7 @@ def init_db():
  db = get_db()
 
  with current_app.open_resource('database.sql') as f:
- db.raw_execute(f.read().decode('utf8'))
+ db.execute_query(f.read().decode('utf8'))
 
 
 @click.command('init-db')

page_analyzer/services/db.py

@@ -1,6 +1,6 @@
 import psycopg
 import logging
-from typing import Optional
+from typing import Optional, List, Tuple, Any
 
 from psycopg.rows import dict_row
 
@@ -11,45 +11,44 @@ def __init__(self, dsn):
  self.logger = logging.getLogger(__name__)
 
  try:
- self.connect = psycopg.connect(dsn,  row_factory=dict_row)
+ self.connection = psycopg.connect(dsn, row_factory=dict_row)
  except Exception as e:
  self.logger.error(f"Ошибка при подключении к базе данных: {str(e)}")
 
  def close(self) -> None:
 
  """Закрывает соединение с базой данных"""
 
- self.connect.close()
+ self.connection.close()
  self.logger.info("Соединение с базой данных закрыто")
 
- def raw_execute(self, query: str) -> None:
- connect = self.connect
- try:
- connect.execute(query)
- connect.commit()
- except Exception as e:
- self.logger.error(f"Ошибка при выполнении запроса: {str(e)}")
- connect.rollback()
+ def execute_query(self,
+ query: str,
+ params: Optional[tuple] = None,
+ commit: bool = False) -> Optional[list[Any]]:
+ """Запрос к бд"""
 
- def execute_query(
- self, query: str,
- params: Optional[tuple] = None,
- commit: bool = False
- ) -> Optional[list]:
- """Выполняет SQL-запрос к базе данных"""
-
- connect = self.connect
+ result = None
+ cursor = self.connection.cursor()
 
  try:
- result = connect.execute(query, params).fetchall()
+ cursor.execute(query, params)
+
+ if cursor.description:
+ result = cursor.fetchall()
+
  if commit:
- connect.commit()
- return result
- except Exception as e:
+ self.connection.commit()
+
+ except psycopg.Error as e:
  self.logger.error(f"Ошибка при выполнении запроса: {str(e)}")
- connect.rollback()
- return None
+
+ finally:
+ if cursor:
+ cursor.close()
+
+ return result
 
  def is_closed(self):
- if self.connect:
- return self.connect.closed
+ if self.connection:
+ return self.connection.closed

tests/conftest.py

@@ -25,7 +25,7 @@ def app():
  with app.app_context():
  db = get_db()
  init_db()
- db.raw_execute(_test_data_sql)
+ db.execute_query(_test_data_sql)
 
  yield app
 

tests/test_db.py

@@ -13,3 +13,11 @@ def test_get_close_db(app):
  db.execute_query('SELECT 1')
 
  assert 'closed' in str(error.value)
+
+
+def test_get_data_db(app):
+ with app.app_context():
+ db = get_db()
+ result = db.execute_query('select * from urls;')
+ assert result[0]['id'] == 1
+ assert result[1]['id'] == 2