feat: migrate batch generator backend to SQL Server with auto-provisioned tables

app.py

@@ -2622,14 +2622,56 @@ def api_content_planner_generate():
         return jsonify({'error': f"Lỗi lập kế hoạch nội dung: {_err_msg(str(e))}"}), 500
 
 # ==============================================
-# TOOL: AI BATCH GENERATOR
+# TOOL: AI BATCH GENERATOR (SQL SERVER CONNECTED)
 # ==============================================
 import uuid
 import threading
+import pymssql
+
+def get_db_connection():
+    return pymssql.connect(
+        server='db57971.public.databaseasp.net',
+        user='db57971',
+        password='Qn4@3F#p!9fD',
+        database='db57971',
+        charset='utf8'
+    )
 
-BATCH_JOBS = {}
+def init_db():
+    try:
+        conn = get_db_connection()
+        cursor = conn.cursor()
+        cursor.execute("""
+        IF NOT EXISTS (SELECT * FROM sysobjects WHERE name='batch_jobs' AND xtype='U')
+        CREATE TABLE batch_jobs (
+            id VARCHAR(50) PRIMARY KEY,
+            mode VARCHAR(10) NOT NULL,
+            status VARCHAR(20) NOT NULL,
+            created_at DATETIME DEFAULT GETDATE()
+        )
+        """)
+        cursor.execute("""
+        IF NOT EXISTS (SELECT * FROM sysobjects WHERE name='batch_prompts' AND xtype='U')
+        CREATE TABLE batch_prompts (
+            id INT IDENTITY(1,1) PRIMARY KEY,
+            batch_id VARCHAR(50) FOREIGN KEY REFERENCES batch_jobs(id) ON DELETE CASCADE,
+            prompt_text NVARCHAR(MAX) NOT NULL,
+            status VARCHAR(20) NOT NULL,
+            output_url VARCHAR(2083) NULL,
+            error_message NVARCHAR(MAX) NULL,
+            created_at DATETIME DEFAULT GETDATE()
+        )
+        """)
+        conn.commit()
+        conn.close()
+        print("SQL Server initialized and tables verified successfully.")
+    except Exception as e:
+        print(f"Error initializing SQL Server database: {e}")
 
-def run_batch_replicate(batch_id, prompts, mode, replicate_token):
+# Run database setup on startup
+init_db()
+
+def run_batch_replicate_db(batch_id, mode, replicate_token):
     # Model configuration
     if mode == 'video':
         model_version = "748dae3f714ccf25a72064119d854203f1917b2b8c56fa769f3e498c0d9c491f"
@@ -2641,10 +2683,31 @@ def run_batch_replicate(batch_id, prompts, mode, replicate_token):
         'Content-Type': 'application/json'
     }
 
-    for idx, prompt_item in enumerate(prompts):
-        prompt_text = prompt_item['prompt']
-        prompt_item['status'] = 'processing'
-        
+    try:
+        # Fetch all prompts to process
+        conn = get_db_connection()
+        cursor = conn.cursor(as_dict=True)
+        cursor.execute("SELECT id, prompt_text FROM batch_prompts WHERE batch_id = %s", (batch_id,))
+        prompts = cursor.fetchall()
+        conn.close()
+    except Exception as e:
+        print(f"Error fetching prompts from DB: {e}")
+        return
+
+    for prompt_item in prompts:
+        prompt_id = prompt_item['id']
+        prompt_text = prompt_item['prompt_text']
+
+        # 1. Update status to 'processing'
+        try:
+            conn = get_db_connection()
+            cursor = conn.cursor()
+            cursor.execute("UPDATE batch_prompts SET status = 'processing' WHERE id = %d", (prompt_id,))
+            conn.commit()
+            conn.close()
+        except Exception as e:
+            print(f"Error updating status to processing: {e}")
+
         # Prepare payload
         payload = {
             "version": model_version,
@@ -2655,7 +2718,7 @@ def run_batch_replicate(batch_id, prompts, mode, replicate_token):
             payload["input"]["guidance_scale"] = 1.0
 
         try:
-            # 1. Create prediction request
+            # Create prediction request
             req = urllib.request.Request(
                 "https://api.replicate.com/v1/predictions",
                 data=json.dumps(payload).encode('utf-8'),
@@ -2669,7 +2732,7 @@ def run_batch_replicate(batch_id, prompts, mode, replicate_token):
             if not prediction_id:
                 raise Exception("Không nhận được prediction_id từ Replicate")
 
-            # 2. Polling loop
+            # Polling loop
             poll_url = f"https://api.replicate.com/v1/predictions/{prediction_id}"
             success = False
             for _ in range(30): # max 60 seconds polling (30 * 2)
@@ -2682,33 +2745,60 @@ def run_batch_replicate(batch_id, prompts, mode, replicate_token):
                 pred_status = status_data.get('status')
                 if pred_status == 'succeeded':
                     output = status_data.get('output')
+                    output_url = None
                     if isinstance(output, list) and len(output) > 0:
-                        prompt_item['output_url'] = output[0]
+                        output_url = output[0]
                     elif isinstance(output, str):
-                        prompt_item['output_url'] = output
+                        output_url = output
                     else:
-                        prompt_item['output_url'] = str(output)
+                        output_url = str(output)
                     
-                    prompt_item['status'] = 'done'
+                    # Update status to 'done' in DB
+                    conn = get_db_connection()
+                    cursor = conn.cursor()
+                    cursor.execute("UPDATE batch_prompts SET status = 'done', output_url = %s WHERE id = %d", (output_url, prompt_id))
+                    conn.commit()
+                    conn.close()
                     success = True
                     break
                 elif pred_status in ['failed', 'canceled']:
-                    prompt_item['status'] = 'failed'
-                    prompt_item['error'] = status_data.get('error') or "Đã bị hủy hoặc lỗi"
+                    error_msg = status_data.get('error') or "Đã bị hủy hoặc lỗi"
+                    # Update status to 'failed' in DB
+                    conn = get_db_connection()
+                    cursor = conn.cursor()
+                    cursor.execute("UPDATE batch_prompts SET status = 'failed', error_message = %s WHERE id = %d", (error_msg, prompt_id))
+                    conn.commit()
+                    conn.close()
                     success = True
                     break
             
             if not success:
-                prompt_item['status'] = 'failed'
-                prompt_item['error'] = "Quá thời gian chờ phản hồi (Timeout)"
+                conn = get_db_connection()
+                cursor = conn.cursor()
+                cursor.execute("UPDATE batch_prompts SET status = 'failed', error_message = 'Quá thời gian chờ phản hồi (Timeout)' WHERE id = %d", (prompt_id,))
+                conn.commit()
+                conn.close()
 
         except Exception as e:
-            prompt_item['status'] = 'failed'
-            prompt_item['error'] = str(e)
+            # Update status to 'failed' in DB
+            try:
+                conn = get_db_connection()
+                cursor = conn.cursor()
+                cursor.execute("UPDATE batch_prompts SET status = 'failed', error_message = %s WHERE id = %d", (str(e), prompt_id))
+                conn.commit()
+                conn.close()
+            except Exception as dbe:
+                print(f"Error updating failure in DB: {dbe}")
 
     # Mark batch job as completed
-    if batch_id in BATCH_JOBS:
-        BATCH_JOBS[batch_id]['status'] = 'completed'
+    try:
+        conn = get_db_connection()
+        cursor = conn.cursor()
+        cursor.execute("UPDATE batch_jobs SET status = 'completed' WHERE id = %s", (batch_id,))
+        conn.commit()
+        conn.close()
+    except Exception as e:
+        print(f"Error updating batch job completion: {e}")
 
 
 @app.route('/batch-generator')
@@ -2729,33 +2819,28 @@ def api_batch_generator_start():
     if not raw_prompts or not isinstance(raw_prompts, list):
         return jsonify({'error': 'Danh sách prompt không hợp lệ'}), 400
 
-    # Build prompts structure
-    prompts_list = []
-    for p in raw_prompts:
-        p_str = str(p).strip()
-        if p_str:
-            prompts_list.append({
-                'prompt': p_str,
-                'status': 'queued',
-                'output_url': None,
-                'error': None
-            })
-
+    prompts_list = [str(p).strip() for p in raw_prompts if str(p).strip()]
     if not prompts_list:
         return jsonify({'error': 'Không tìm thấy prompt hợp lệ nào'}), 400
 
     batch_id = str(uuid.uuid4())
-    BATCH_JOBS[batch_id] = {
-        'id': batch_id,
-        'mode': mode,
-        'status': 'running',
-        'prompts': prompts_list
-    }
+
+    try:
+        # Insert batch job and prompts in database
+        conn = get_db_connection()
+        cursor = conn.cursor()
+        cursor.execute("INSERT INTO batch_jobs (id, mode, status) VALUES (%s, %s, %s)", (batch_id, mode, 'running'))
+        for prompt_text in prompts_list:
+            cursor.execute("INSERT INTO batch_prompts (batch_id, prompt_text, status) VALUES (%s, %s, %s)", (batch_id, prompt_text, 'queued'))
+        conn.commit()
+        conn.close()
+    except Exception as e:
+        return jsonify({'error': f'Lỗi lưu thông tin vào CSDL SQL Server: {e}'}), 500
 
     # Start processing thread
     thread = threading.Thread(
-        target=run_batch_replicate,
-        args=(batch_id, prompts_list, mode, replicate_token)
+        target=run_batch_replicate_db,
+        args=(batch_id, mode, replicate_token)
     )
     thread.daemon = True
     thread.start()
@@ -2769,15 +2854,27 @@ def api_batch_generator_start():
 
 @app.route('/api/batch-generator/status/<batch_id>')
 def api_batch_generator_status(batch_id):
-    job = BATCH_JOBS.get(batch_id)
-    if not job:
-        return jsonify({'error': 'Không tìm thấy mã Batch'}), 404
-    return jsonify({
-        'success': True,
-        'status': job['status'],
-        'mode': job['mode'],
-        'prompts': job['prompts']
-    })
+    try:
+        conn = get_db_connection()
+        cursor = conn.cursor(as_dict=True)
+        cursor.execute("SELECT status, mode FROM batch_jobs WHERE id = %s", (batch_id,))
+        job = cursor.fetchone()
+        if not job:
+            conn.close()
+            return jsonify({'error': 'Không tìm thấy mã Batch'}), 404
+        
+        cursor.execute("SELECT prompt_text AS prompt, status, output_url, error_message AS error FROM batch_prompts WHERE batch_id = %s ORDER BY id ASC", (batch_id,))
+        prompts = cursor.fetchall()
+        conn.close()
+
+        return jsonify({
+            'success': True,
+            'status': job['status'],
+            'mode': job['mode'],
+            'prompts': prompts
+        })
+    except Exception as e:
+        return jsonify({'error': f'Lỗi truy vấn CSDL SQL Server: {e}'}), 500
 
 if __name__ == '__main__':
     app.run(debug=True, port=5000)

requirements.txt

@@ -14,3 +14,4 @@ vieneu
 moviepy
 replicate
 gradio-client
+pymssql