Skip to main content

sqlite_knowledge_graph/
schema.rs

1//! Database schema creation and management.
2
3use rusqlite::Connection;
4
5use crate::error::Result;
6
7/// Create the knowledge graph schema in the database.
8pub fn create_schema(conn: &Connection) -> Result<()> {
9    let tx = conn.unchecked_transaction()?;
10
11    // Create entities table
12    tx.execute(
13        r#"
14        CREATE TABLE IF NOT EXISTS kg_entities (
15            id INTEGER PRIMARY KEY AUTOINCREMENT,
16            entity_type TEXT NOT NULL,
17            name TEXT NOT NULL,
18            properties TEXT,  -- JSON
19            created_at INTEGER DEFAULT (strftime('%s', 'now')),
20            updated_at INTEGER DEFAULT (strftime('%s', 'now'))
21        )
22        "#,
23        [],
24    )?;
25
26    // Create relations table
27    tx.execute(
28        r#"
29        CREATE TABLE IF NOT EXISTS kg_relations (
30            id INTEGER PRIMARY KEY AUTOINCREMENT,
31            source_id INTEGER NOT NULL,
32            target_id INTEGER NOT NULL,
33            rel_type TEXT NOT NULL,
34            weight REAL DEFAULT 1.0,
35            properties TEXT,  -- JSON
36            created_at INTEGER DEFAULT (strftime('%s', 'now')),
37            FOREIGN KEY (source_id) REFERENCES kg_entities(id) ON DELETE CASCADE,
38            FOREIGN KEY (target_id) REFERENCES kg_entities(id) ON DELETE CASCADE
39        )
40        "#,
41        [],
42    )?;
43
44    // Create vectors table
45    tx.execute(
46        r#"
47        CREATE TABLE IF NOT EXISTS kg_vectors (
48            entity_id INTEGER NOT NULL PRIMARY KEY,
49            vector BLOB NOT NULL,
50            dimension INTEGER NOT NULL,
51            created_at INTEGER DEFAULT (strftime('%s', 'now')),
52            FOREIGN KEY (entity_id) REFERENCES kg_entities(id) ON DELETE CASCADE
53        )
54        "#,
55        [],
56    )?;
57
58    // Create indexes
59    tx.execute(
60        "CREATE INDEX IF NOT EXISTS idx_entities_type ON kg_entities(entity_type)",
61        [],
62    )?;
63
64    tx.execute(
65        "CREATE INDEX IF NOT EXISTS idx_entities_name ON kg_entities(name)",
66        [],
67    )?;
68
69    tx.execute(
70        "CREATE INDEX IF NOT EXISTS idx_relations_source ON kg_relations(source_id)",
71        [],
72    )?;
73
74    tx.execute(
75        "CREATE INDEX IF NOT EXISTS idx_relations_target ON kg_relations(target_id)",
76        [],
77    )?;
78
79    tx.execute(
80        "CREATE INDEX IF NOT EXISTS idx_relations_type ON kg_relations(rel_type)",
81        [],
82    )?;
83
84    tx.commit()?;
85    Ok(())
86}
87
88/// Check if the schema exists.
89pub fn schema_exists(conn: &Connection) -> Result<bool> {
90    let mut stmt = conn
91        .prepare("SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='kg_entities'")?;
92    let count: i64 = stmt.query_row([], |row| row.get(0))?;
93    Ok(count > 0)
94}