init database expect db_node

ntfs_utils/db_device.py

@@ -0,0 +1,125 @@
+import sqlite3
+
+import psutil
+
+
+def ScanSpecialVolumes(volume_letter):
+    """
+    扫描指定的单个磁盘卷，返回其基本信息字典。
+
+    参数:
+        volume_letter: str，磁盘盘符（例如 "C", "Z"）
+
+    返回:
+        dict: 包含 Path, Type, Option 字段的字典
+    """
+    # 简单校验盘符格式（去除可能的冒号）
+    if len(volume_letter) >= 1 and volume_letter[-1] == ":":
+        volume_letter = volume_letter[:-1]
+
+    if not volume_letter or len(volume_letter) != 1 or not volume_letter.isalpha():
+        raise ValueError("无效的磁盘盘符，应为单个字母，如 'C' 或 'Z'")
+
+    return {
+        "Path": volume_letter.upper(),
+        "Type": "磁盘",
+        "Option": None
+    }
+
+
+def ScanNTFSVolumes():
+    """
+    扫描当前系统中所有 NTFS 格式的磁盘卷。
+
+    返回:
+        list of dict: 包含盘符等信息的字典列表
+    """
+    ntfs_volumes = []
+
+    for partition in psutil.disk_partitions():
+        if partition.fstype.upper() == 'NTFS':
+            # 提取盘符（去掉冒号）
+            drive_letter = partition.device[0] if len(partition.device) >= 2 and partition.device[1] == ':' else None
+            if drive_letter:
+                ntfs_volumes.append({
+                    "Path": drive_letter,
+                    "Type": "磁盘",
+                    "Option": None
+                })
+
+    return ntfs_volumes
+
+
+def InsertVolumesToDB(data, db_path='../src/db_ntfs_info.db', table_name='db_device'):
+    """
+    将 NTFS 磁盘信息写入数据库表，并防止重复插入。
+
+    参数:
+        data: list of dict，包含 Path, Type, Option 字段的数据
+        db_path: str，SQLite 数据库路径
+        table_name: str，目标表名
+
+    返回:
+        int: 成功插入的记录数
+    """
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+
+    try:
+        # 创建表（如果不存在），并添加 Path 的唯一性约束
+        create_table_sql = f"""
+        CREATE TABLE IF NOT EXISTS {table_name} (
+            ID INTEGER PRIMARY KEY AUTOINCREMENT,
+            Path TEXT NOT NULL UNIQUE,
+            Type TEXT NOT NULL CHECK(Type IN ('磁盘', '文件', '文件夹')),
+            Option TEXT
+        );
+        """
+        cursor.execute(create_table_sql)
+
+        inserted_count = 0
+        insert_sql = f"""
+        INSERT OR IGNORE INTO {table_name} (Path, Type, Option)
+        VALUES (?, ?, ?)
+        """
+
+        for item in data:
+            cursor.execute(insert_sql, (
+                item['Path'],
+                item['Type'],
+                item['Option']
+            ))
+            if cursor.rowcount > 0:
+                inserted_count += 1
+
+        conn.commit()
+        print(f"✅ 成功插入 {inserted_count} 条 NTFS 磁盘信息")
+        return inserted_count
+
+    except Exception as e:
+        print(f"❌ 数据库操作失败: {e}")
+        conn.rollback()
+        return 0
+
+    finally:
+        conn.close()
+
+
+def main():
+    # 扫描系统下所有 NTFS 磁盘
+    # volumes = ScanNTFSVolumes()
+    # print("🔍 找到以下 NTFS 磁盘：")
+    # for vol in volumes:
+    #     print(vol)
+    #
+    # success_count = InsertVolumesToDB(volumes)
+    # print(f"共插入 {success_count} 条记录到数据库。")
+
+    # 扫描单个磁盘
+    volume_latter = "Z"
+    device_data = ScanSpecialVolumes(volume_latter)
+    InsertVolumesToDB([device_data])
+
+
+if __name__ == "__main__":
+    main()