추출횟수 업데이트

main.py

@@ -347,6 +347,24 @@ class MainWindow(QMainWindow):
         self.run_button.setToolTip("선택된 필터 조건에 따라 DB에서 데이터를 가져와 북마크를 추가합니다.")
         self.run_button.clicked.connect(self.run_task)
 
+        # 추출 옵션 체크박스
+        self.extract_based_checkbox = QCheckBox("추출 횟수 기반 추출")
+        self.extract_based_checkbox.setChecked(False)  # 기본값 OFF
+        self.extract_based_checkbox.setToolTip("체크 시 추출 횟수가 낮은 데이터부터 추출합니다.")
+        self.extract_based_checkbox.stateChanged.connect(self.toggle_extract_options)
+
+        # 최대 추출 횟수 스핀박스
+        self.max_extract_label = QLabel("최대 추출 횟수")
+        self.max_extract_spinbox = QSpinBox()
+        self.max_extract_spinbox.setMinimum(1)
+        self.max_extract_spinbox.setMaximum(5)
+        self.max_extract_spinbox.setValue(1)  # 기본값
+
+        # 리셋 버튼
+        self.reset_button = QPushButton("추출 횟수 초기화")
+        self.reset_button.setToolTip("DB의 모든 추출 횟수를 0으로 초기화합니다.")
+        self.reset_button.clicked.connect(self.reset_extract_count)
+
         # 로그 박스
         self.log_box = QTextEdit()
         self.log_box.setReadOnly(True)
@@ -384,6 +402,11 @@ class MainWindow(QMainWindow):
         self.buttons_layout.addWidget(self.view_data_button)
         self.buttons_layout.addWidget(self.run_button)
 
+        self.filter_layout.addWidget(self.extract_based_checkbox)
+        self.filter_layout.addWidget(self.max_extract_label)
+        self.filter_layout.addWidget(self.max_extract_spinbox)
+        self.filter_layout.addWidget(self.reset_button)
+
         # 메인 레이아웃 구성
         self.layout.addLayout(self.filter_layout)
         self.layout.addLayout(self.buttons_layout)
@@ -648,11 +671,19 @@ class MainWindow(QMainWindow):
                         country TEXT,
                         mall_grade TEXT,
                         mall_name TEXT,
-                        mall_url TEXT
+                        mall_url TEXT,
+                        extract_count INTEGER DEFAULT 0
                     )
                 """)
                 self.log("기존 DB가 제거되었습니다. 새롭게 생성되었습니다.")
 
+            # 기존 DB에 `extract_count` 필드 추가
+            try:
+                conn.execute("ALTER TABLE markets ADD COLUMN extract_count INTEGER DEFAULT 0")
+            except sqlite3.OperationalError:
+                # 컬럼이 이미 있는 경우 무시
+                pass
+
             # 데이터 저장, 중복 방지
             df.drop_duplicates(subset=['country', 'mall_grade', 'mall_name', 'mall_url'], inplace=True)
 
@@ -756,38 +787,93 @@ class MainWindow(QMainWindow):
         else:
             self.log("수정이 취소되었습니다.")
 
+    def ensure_extract_count_column(self, conn):
+        cursor = conn.cursor()
+        cursor.execute("PRAGMA table_info(markets)")
+        columns = [row[1] for row in cursor.fetchall()]
+        if "extract_count" not in columns:
+            conn.execute("ALTER TABLE markets ADD COLUMN extract_count INTEGER DEFAULT 0")
+
+    def toggle_extract_options(self):
+        if self.extract_based_checkbox.isChecked():
+            self.max_extract_spinbox.setEnabled(True)
+        else:
+            self.max_extract_spinbox.setEnabled(False)
+
     def run_task(self):
         country = self.country_dropdown.currentText()
         grade = self.grade_dropdown.currentText()
         count = self.count_spinbox.value()
         remove_existing = self.remove_existing_checkbox.isChecked()
 
+        extract_based = self.extract_based_checkbox.isChecked()
+        max_extract = self.max_extract_spinbox.value()
+
         conn = sqlite3.connect(self.db_path)
-        query = "SELECT mall_name AS name, mall_url AS url FROM markets WHERE 1=1"
-        
-        # 국가 필터
-        if country != "랜덤":
-            query += f" AND country = '{country}'"
-        # 등급 필터
-        if grade != "랜덤":
-            query += f" AND mall_grade = '{grade}'"
-        
-        # 랜덤 정렬 추가
-        query += " ORDER BY RANDOM()"
-        # 갯수 제한 추가
-        query += f" LIMIT {count}"
 
         try:
-            df = pd.read_sql_query(query, conn)
-            conn.close()
+            # 추출 횟수 필드 확인 및 추가
+            self.ensure_extract_count_column(conn)
+
+            # 기본 쿼리
+            query = "SELECT id, mall_name AS name, mall_url AS url FROM markets WHERE 1=1"
+            
+            # 국가 필터
+            if country != "랜덤":
+                query += f" AND country = '{country}'"
+            # 등급 필터
+            if grade != "랜덤":
+                query += f" AND mall_grade = '{grade}'"
+            
+            # 추출 횟수 기반 옵션 처리
+            if extract_based:
+                query += " AND extract_count < ?"
+                query += " ORDER BY extract_count ASC, RANDOM()"
+            else:
+                query += " ORDER BY RANDOM()"
+            
+            query += f" LIMIT {count}"
+
+            # 추출
+            if extract_based:
+                df = pd.read_sql_query(query, conn, params=(max_extract,))
+            else:
+                df = pd.read_sql_query(query, conn)
+
+            if df.empty:
+                # 추출 가능한 데이터가 없는 경우 리셋
+                if extract_based:
+                    reply = QMessageBox.question(
+                        self,
+                        "데이터 초기화",
+                        "추출 가능한 데이터가 없습니다. 모든 추출 횟수를 초기화하시겠습니까?",
+                        QMessageBox.Yes | QMessageBox.No
+                    )
+                    if reply == QMessageBox.Yes:
+                        self.reset_extract_count()
+                    else:
+                        return
+                else:
+                    QMessageBox.warning(self, "추출 실패", "DB에 추출 가능한 데이터가 없습니다.")
+                    return
 
             self.bookmarks = df.to_dict("records")
+
+            # 추출 횟수 업데이트
+            for record in df.to_dict("records"):
+                conn.execute("UPDATE markets SET extract_count = extract_count + 1 WHERE id = ?", (record["id"],))
+            conn.commit()
+
             self.log(f"{len(self.bookmarks)}개의 북마크를 추출했습니다.")
+
         except Exception as e:
             self.log(f"DB 쿼리 실행 중 오류 발생: {str(e)}", exc_info=True)
             conn.close()
             return
 
+        conn.close()
+
+
         folder_name = f"거상북마크-{grade}"
 
         # 사용자가 선택한 브라우저에 따라 경로 설정
@@ -825,6 +911,25 @@ class MainWindow(QMainWindow):
         self.log("작업이 완료되었습니다!")
         QMessageBox.information(self, "완료", "즐겨찾기 추가 작업이 완료되었습니다.")
 
+    def reset_extract_count(self):
+        conn = sqlite3.connect(self.db_path)
+        try:
+            # `extract_count` 필드가 없으면 추가
+            try:
+                conn.execute("ALTER TABLE markets ADD COLUMN extract_count INTEGER DEFAULT 0")
+            except sqlite3.OperationalError:
+                # 이미 필드가 존재하는 경우 무시
+                pass
+
+            # 모든 추출 횟수를 0으로 초기화
+            conn.execute("UPDATE markets SET extract_count = 0")
+            conn.commit()
+            self.log("모든 추출 횟수가 초기화되었습니다.")
+            QMessageBox.information(self, "초기화 완료", "모든 추출 횟수가 초기화되었습니다.")
+        except Exception as e:
+            self.log(f"추출 횟수 초기화 중 오류 발생: {str(e)}", exc_info=True)
+        finally:
+            conn.close()
 
 
 def main():