baidu_web/databaseManager.py

153 lines
7.0 KiB
Python

import sqlite3
import pandas as pd
class DatabaseManager:
def __init__(self, db_path, logger):
self.logger = logger
self.conn = sqlite3.connect(db_path)
self.logger.info("Database connection established.")
self._create_tables()
def _create_tables(self):
try:
with self.conn:
# products 테이블 생성
self.conn.execute('''
CREATE TABLE IF NOT EXISTS products (
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT,
image_url TEXT,
tag TEXT,
price INTEGER,
category TEXT,
percenty_category TEXT,
selected_search_result_id INTEGER,
saved_img_path TEXT, -- 이미지 저장 경로 필드 추가
FOREIGN KEY(selected_search_result_id) REFERENCES search_results(id)
)
''')
# search_results 테이블에 고유 ID 추가
self.conn.execute('''
CREATE TABLE IF NOT EXISTS search_results (
id INTEGER PRIMARY KEY AUTOINCREMENT,
product_id INTEGER,
title TEXT,
source TEXT,
price TEXT,
imgurl TEXT,
encrypted_url TEXT,
original_url TEXT,
saved_img_path TEXT, -- 이미지 저장 경로 필드 추가
is_selected BOOLEAN DEFAULT 0,
FOREIGN KEY (product_id) REFERENCES products (id)
)
''')
self.logger.info("Tables created successfully.")
except Exception as e:
self.logger.error(f"Error creating tables: {e}", exc_info=True)
def insert_products(self, products):
# 엑셀에서 가져온 상품 정보를 DB에 개별적으로 저장하고 product_id 리스트 반환
product_ids = []
try:
with self.conn:
for product in products:
cursor = self.conn.execute('''
INSERT INTO products (name, image_url, tag, price, category, percenty_category, saved_img_path)
VALUES (?, ?, ?, ?, ?, ?, ?)
''', (product['name'], product['image_url'], product['tag'], product['price'], product['category'], product['percenty_category'], product.get('saved_img_path')))
product_id = cursor.lastrowid
product_ids.append(product_id)
self.logger.debug(f"Inserted product with ID: {product_id}")
self.logger.info(f"{len(products)} products inserted into database.")
return product_ids
except Exception as e:
self.logger.error(f"Error inserting products: {e}", exc_info=True)
return []
def update_product_image_path(self, product_id, image_path):
"""특정 상품의 saved_img_path 필드만 업데이트"""
try:
with self.conn:
self.conn.execute('''
UPDATE products
SET saved_img_path = ?
WHERE id = ?
''', (image_path, product_id))
self.logger.info(f"Updated saved_img_path for Product ID [{product_id}]")
except Exception as e:
self.logger.error(f"Error updating saved_img_path for Product ID [{product_id}]: {e}", exc_info=True)
def insert_search_results(self, product_id, search_results):
# 검색 결과 중 상위 5개만 DB에 저장
try:
top_results = search_results[:5] # 상위 5개로 제한
with self.conn:
self.conn.executemany('''
INSERT INTO search_results (product_id, title, source, price, imgurl, saved_img_path, encrypted_url, original_url)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
''', [(product_id, result['title'], result['source'], result['price'], result['imgurl'], result['saved_img_path'], result['encrypted_url'], result['original_url']) for result in top_results])
self.logger.info(f"Product ID [{product_id}]: {len(top_results)} search results inserted into database.")
except Exception as e:
self.logger.error(f"Error inserting search results for Product ID [{product_id}]: {e}", exc_info=True)
def select_search_result_for_product(self, product_id, search_result_id):
try:
with self.conn:
# products 테이블의 selected_search_result_id 업데이트
self.conn.execute('''
UPDATE products
SET selected_search_result_id = ?
WHERE id = ?
''', (search_result_id, product_id))
# search_results 테이블의 선택된 항목 업데이트
self.conn.execute('''
UPDATE search_results
SET is_selected = CASE WHEN id = ? THEN 1 ELSE 0 END
WHERE product_id = ?
''', (search_result_id, product_id))
self.logger.info(f"Product ID [{product_id}] selected search result ID [{search_result_id}].")
except Exception as e:
self.logger.error(f"Error selecting search result for product ID [{product_id}]: {e}", exc_info=True)
def get_selected_search_result(self, product_id):
# products 테이블에서 선택된 search_result_id 가져오기
try:
cursor = self.conn.execute('''
SELECT selected_search_result_id
FROM products
WHERE id = ?
''', (product_id,))
row = cursor.fetchone()
if row and row[0]:
return row[0]
return None
except Exception as e:
self.logger.error(f"Error fetching selected search result for product ID [{product_id}]: {e}", exc_info=True)
return None
def close(self):
if self.conn:
self.conn.close()
self.logger.info("Database connection closed.")
def export_to_excel(self, excel_path):
try:
query = '''
SELECT p.name as 상품명, sr.original_url as 원본링크, sr.price as 가격,
sr.source as 출처, p.percenty_category as 퍼센티카테고리,
sr.title as 이미지검색결과의상품명, sr.imgurl as 이미지URL
FROM products p
LEFT JOIN search_results sr ON p.id = sr.product_id
'''
df = pd.read_sql_query(query, self.conn)
df.to_excel(excel_path, index=False)
self.logger.info("Data exported to Excel successfully.")
except Exception as e:
self.logger.error(f"Error exporting to Excel: {e}")