Team Ai
Apppublic

Quan1811/docstring-api

sourceHugging Facemitupdated 5mo agoView on Hugging Face
0likes
setup_db.py101 linesDownload Raw Back to root
1import mysql.connector2import os3from dotenv import load_dotenv4import auth  # Import auth để băm mật khẩu5 6load_dotenv()7 8# Cấu hình từ .env9DB_HOST = os.getenv("DB_HOST", "localhost")10DB_USER = os.getenv("DB_USER", "root")11DB_PASS = os.getenv("DB_PASS", "")12DB_NAME = os.getenv("DB_NAME", "docstring_db")13 14def setup():15    try:16        # 1. Kết nối không chọn DB để tạo DB trước17        conn = mysql.connector.connect(18            host=DB_HOST,19            user=DB_USER,20            password=DB_PASS21        )22        cursor = conn.cursor()23        24        # Initialize database25        cursor.execute(f"CREATE DATABASE IF NOT EXISTS {DB_NAME} CHARACTER SET utf8mb4")26        print(f"Database {DB_NAME} initialized.")27        conn.close()28 29        # 2. Kết nối vào DB để tạo bảng30        conn = mysql.connector.connect(31            host=DB_HOST,32            user=DB_USER,33            password=DB_PASS,34            database=DB_NAME35        )36        cursor = conn.cursor()37 38        # Bảng Users39        cursor.execute("""40            CREATE TABLE IF NOT EXISTS users (41                id INT AUTO_INCREMENT PRIMARY KEY,42                username VARCHAR(50) UNIQUE NOT NULL,43                email VARCHAR(100) UNIQUE NOT NULL,44                password VARCHAR(255) NOT NULL45            )46        """)47        48        # Bảng Codes49        cursor.execute("""50            CREATE TABLE IF NOT EXISTS codes (51                id INT AUTO_INCREMENT PRIMARY KEY,52                content TEXT NOT NULL,53                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP54            )55        """)56 57        # Bảng Docstrings58        cursor.execute("""59            CREATE TABLE IF NOT EXISTS docstrings (60                id INT AUTO_INCREMENT PRIMARY KEY,61                content TEXT NOT NULL,62                created_at DATETIME DEFAULT CURRENT_TIMESTAMP63            )64        """)65 66        # Bảng History (Lưu liên kết)67        cursor.execute("""68            CREATE TABLE IF NOT EXISTS history (69                id INT AUTO_INCREMENT PRIMARY KEY,70                user_id INT,71                code_id INT,72                docstring_id INT,73                created_at DATETIME DEFAULT CURRENT_TIMESTAMP,74                FOREIGN KEY (user_id) REFERENCES users(id),75                FOREIGN KEY (code_id) REFERENCES codes(id),76                FOREIGN KEY (docstring_id) REFERENCES docstrings(id)77            )78        """)79        80        print("Tables initialized: users, codes, docstrings, history.")81 82        # Seed data83        cursor.execute("SELECT COUNT(*) FROM users WHERE id = 1")84        if cursor.fetchone()[0] == 0:85            hashed_pw = auth.hash_password("password123")86            cursor.execute(87                "INSERT INTO users (id, username, email, password) VALUES (%s, %s, %s, %s)",88                (1, "testuser", "test@example.com", hashed_pw)89            )90            conn.commit()91            print("Default user created (username: testuser, password: password123).")92 93        conn.close()94        print("Setup completed.")95 96    except Exception as e:97        print(f"Error during setup: {e}")98 99if __name__ == "__main__":100    setup()101