Propagate info about hidden columns

core/schema.rs

@@ -667,6 +667,7 @@ fn create_table(
                     default,
                     unique,
                     collation,
+                    hidden: false,
                 });
             }
             if options.contains(TableOptions::WITHOUT_ROWID) {
@@ -738,6 +739,7 @@ pub struct Column {
     pub default: Option<Expr>,
     pub unique: bool,
     pub collation: Option<CollationSeq>,
+    pub hidden: bool,
 }
 
 impl Column {
@@ -816,6 +818,7 @@ impl From<ColumnDefinition> for Column {
             is_rowid_alias: primary_key && matches!(ty, Type::Integer),
             unique,
             collation,
+            hidden: false,
         }
     }
 }
@@ -1031,6 +1034,7 @@ pub fn sqlite_schema_table() -> BTreeTable {
                 default: None,
                 unique: false,
                 collation: None,
+                hidden: false,
             },
             Column {
                 name: Some("name".to_string()),
@@ -1042,6 +1046,7 @@ pub fn sqlite_schema_table() -> BTreeTable {
                 default: None,
                 unique: false,
                 collation: None,
+                hidden: false,
             },
             Column {
                 name: Some("tbl_name".to_string()),
@@ -1053,6 +1058,7 @@ pub fn sqlite_schema_table() -> BTreeTable {
                 default: None,
                 unique: false,
                 collation: None,
+                hidden: false,
             },
             Column {
                 name: Some("rootpage".to_string()),
@@ -1064,6 +1070,7 @@ pub fn sqlite_schema_table() -> BTreeTable {
                 default: None,
                 unique: false,
                 collation: None,
+                hidden: false,
             },
             Column {
                 name: Some("sql".to_string()),
@@ -1075,6 +1082,7 @@ pub fn sqlite_schema_table() -> BTreeTable {
                 default: None,
                 unique: false,
                 collation: None,
+                hidden: false,
             },
         ],
         unique_sets: None,
@@ -1688,6 +1696,7 @@ mod tests {
                 default: None,
                 unique: false,
                 collation: None,
+                hidden: false,
             }],
             unique_sets: None,
         };

core/translate/insert.rs

@@ -645,25 +645,30 @@ fn resolve_columns_for_insert<'a>(
     let table_columns = table.columns();
     // Case 1: No columns specified - map values to columns in order
     if columns.is_none() {
-        if num_values != table_columns.len() {
+        let mut value_idx = 0;
+        let mut column_mappings = Vec::with_capacity(table_columns.len());
+        for col in table_columns {
+            let mapping = ColumnMapping {
+                column: col,
+                value_index: if col.hidden { None } else { Some(value_idx) },
+                default_value: col.default.as_ref(),
+            };
+            if !col.hidden {
+                value_idx += 1;
+            }
+            column_mappings.push(mapping);
+        }
+
+        if num_values != value_idx {
             crate::bail_parse_error!(
                 "table {} has {} columns but {} values were supplied",
                 &table.get_name(),
-                table_columns.len(),
+                value_idx,
                 num_values
             );
         }
 
-        // Map each column to either its corresponding value index or None
+        return Ok(column_mappings);
-        return Ok(table_columns
-            .iter()
-            .enumerate()
-            .map(|(i, col)| ColumnMapping {
-                column: col,
-                value_index: if i < num_values { Some(i) } else { None },
-                default_value: col.default.as_ref(),
-            })
-            .collect());
     }
 
     // Case 2: Columns specified - map named columns to their values
@@ -852,6 +857,12 @@ fn populate_column_registers(
             if write_directly_to_rowid_reg {
                 program.emit_insn(Insn::SoftNull { reg: target_reg });
             }
+        } else if mapping.column.hidden {
+            program.emit_insn(Insn::Null {
+                dest: target_reg,
+                dest_end: None,
+            });
+            program.mark_last_insn_constant();
         } else if let Some(default_expr) = mapping.default_value {
             translate_expr_no_constant_opt(
                 program,

core/translate/optimizer/join.rs

@@ -1610,6 +1610,7 @@ mod tests {
             default: None,
             unique: false,
             collation: None,
+            hidden: false,
         }
     }
     fn _create_column_of_type(name: &str, ty: Type) -> Column {

core/translate/plan.rs

@@ -555,6 +555,7 @@ pub fn select_star(tables: &[JoinedTable], out_columns: &mut Vec<ResultSetColumn
                 .columns()
                 .iter()
                 .enumerate()
+                .filter(|(_, col)| !col.hidden)
                 .filter(|(_, col)| {
                     // If we are joining with USING, we need to deduplicate the columns from the right table
                     // that are also present in the USING clause.
@@ -917,6 +918,7 @@ impl JoinedTable {
                 default: None,
                 unique: false,
                 collation: None, // FIXME: infer collation from subquery
+                hidden: false,
             })
             .collect();
 

core/translate/planner.rs

@@ -730,17 +730,17 @@ fn parse_join(
         assert!(table_references.joined_tables().len() >= 2);
         let rightmost_table = table_references.joined_tables().last().unwrap();
         // NATURAL JOIN is first transformed into a USING join with the common columns
-        let right_cols = rightmost_table.columns();
         let mut distinct_names: Option<ast::DistinctNames> = None;
         // TODO: O(n^2) maybe not great for large tables or big multiway joins
-        for right_col in right_cols.iter() {
+        // SQLite doesn't use HIDDEN columns for NATURAL joins: https://www3.sqlite.org/src/info/ab09ef427181130b
+        for right_col in rightmost_table.columns().iter().filter(|col| !col.hidden) {
             let mut found_match = false;
             for left_table in table_references
                 .joined_tables()
                 .iter()
                 .take(table_references.joined_tables().len() - 1)
             {
-                for left_col in left_table.columns().iter() {
+                for left_col in left_table.columns().iter().filter(|col| !col.hidden) {
                     if left_col.name == right_col.name {
                         if let Some(distinct_names) = distinct_names.as_mut() {
                             distinct_names
@@ -805,6 +805,7 @@ fn parse_join(
                             .columns()
                             .iter()
                             .enumerate()
+                            .filter(|(_, col)| !natural || !col.hidden)
                             .find(|(_, col)| {
                                 col.name
                                     .as_ref()

core/translate/pragma.rs

@@ -303,7 +303,10 @@ fn query_pragma(
             let base_reg = register;
             program.alloc_registers(5);
             if let Some(table) = table {
-                for (i, column) in table.columns().iter().enumerate() {
+                // According to the SQLite documentation: "The 'cid' column should not be taken to
+                // mean more than 'rank within the current result set'."
+                // Therefore, we enumerate only after filtering out hidden columns.
+                for (i, column) in table.columns().iter().filter(|col| !col.hidden).enumerate() {
                     // cid
                     program.emit_int(i as i64, base_reg);
                     // name

core/translate/schema.rs

@@ -783,6 +783,7 @@ pub fn translate_drop_table(
                 default: None,
                 unique: false,
                 collation: None,
+                hidden: false,
             }],
             is_strict: false,
             unique_sets: None,

core/translate/select.rs

@@ -234,14 +234,14 @@ fn prepare_one_select_plan(
                         ResultColumn::Star => table_references
                             .joined_tables()
                             .iter()
-                            .map(|t| t.columns().len())
+                            .map(|t| t.columns().iter().filter(|col| !col.hidden).count())
                             .sum(),
                         // Guess 5 columns if we can't find the table using the identifier (maybe it's in [brackets] or `tick_quotes`, or miXeDcAse)
                         ResultColumn::TableStar(n) => table_references
                             .joined_tables()
                             .iter()
                             .find(|t| t.identifier == n.0)
-                            .map(|t| t.columns().len())
+                            .map(|t| t.columns().iter().filter(|col| !col.hidden).count())
                             .unwrap_or(5),
                         // Otherwise allocate space for 1 column
                         ResultColumn::Expr(_, _) => 1,
@@ -284,6 +284,10 @@ fn prepare_one_select_plan(
                         );
                         for table in plan.table_references.joined_tables_mut() {
                             for idx in 0..table.columns().len() {
+                                let column = &table.columns()[idx];
+                                if column.hidden {
+                                    continue;
+                                }
                                 table.mark_column_used(idx);
                             }
                         }
@@ -302,16 +306,16 @@ fn prepare_one_select_plan(
                         let table = referenced_table.unwrap();
                         let num_columns = table.columns().len();
                         for idx in 0..num_columns {
-                            let is_rowid_alias = {
+                            let column = &table.columns()[idx];
-                                let columns = table.columns();
+                            if column.hidden {
-                                columns[idx].is_rowid_alias
+                                continue;
-                            };
+                            }
                             plan.result_columns.push(ResultSetColumn {
                                 expr: ast::Expr::Column {
                                     database: None, // TODO: support different databases
                                     table: table.internal_id,
                                     column: idx,
-                                    is_rowid_alias,
+                                    is_rowid_alias: column.is_rowid_alias,
                                 },
                                 alias: None,
                                 contains_aggregates: false,

core/translate/update.rs

@@ -250,6 +250,7 @@ pub fn prepare_update_plan(
                 default: None,
                 unique: false,
                 collation: None,
+                hidden: false,
             }],
             is_strict: false,
             unique_sets: None,

core/util.rs

@@ -490,14 +490,7 @@ pub fn columns_from_create_table_body(body: &ast::CreateTableBody) -> crate::Res
 
     Ok(columns
         .into_iter()
-        .filter_map(|(name, column_def)| {
+        .map(|(name, column_def)| {
-            // if column_def.col_type includes HIDDEN, omit it for now
-            if let Some(data_type) = column_def.col_type.as_ref() {
-                if data_type.name.as_str().contains("HIDDEN") {
-                    return None;
-                }
-            }
-            let column =
             Column {
                 name: Some(normalize_ident(&name.0)),
                 ty: match column_def.col_type {
@@ -566,15 +559,21 @@ pub fn columns_from_create_table_body(body: &ast::CreateTableBody) -> crate::Res
                         // But in the future, when a user defines a collation sequence, creates a table with it,
                         // then closes the db and opens it again. This may panic here if the collation seq is not registered
                         // before reading the columns
-                            turso_sqlite3_parser::ast::ColumnConstraint::Collate {
+                        turso_sqlite3_parser::ast::ColumnConstraint::Collate { collation_name } => {
-                                collation_name,
+                            Some(
-                            } => Some(CollationSeq::new(collation_name.0.as_str()).expect(
+                                CollationSeq::new(collation_name.0.as_str()).expect(
                                     "collation should have been set correctly in create table",
-                            )),
+                                ),
+                            )
+                        }
                         _ => None,
                     }),
-                };
+                hidden: column_def
-            Some(column)
+                    .col_type
+                    .as_ref()
+                    .map(|data_type| data_type.name.as_str().contains("HIDDEN"))
+                    .unwrap_or(false),
+            }
         })
         .collect::<Vec<_>>())
 }

extensions/tests/src/lib.rs

@@ -19,14 +19,14 @@ register_extension! {
     vfs: { TestFS },
 }
 
-type Store = Rc<RefCell<BTreeMap<i64, (String, String)>>>;
+type Store = Rc<RefCell<BTreeMap<i64, (String, String, String)>>>;
 
 #[derive(VTabModuleDerive, Default)]
 pub struct KVStoreVTabModule;
 
-/// the cursor holds a snapshot of (rowid, key, value) in memory.
+/// the cursor holds a snapshot of (rowid, comment, key, value) in memory.
 pub struct KVStoreCursor {
-    rows: Vec<(i64, String, String)>,
+    rows: Vec<(i64, String, String, String)>,
     index: Option<usize>,
     store: Store,
 }
@@ -37,7 +37,17 @@ impl VTabModule for KVStoreVTabModule {
     const NAME: &'static str = "kv_store";
 
     fn create(_args: &[Value]) -> Result<(String, Self::Table), ResultCode> {
-        let schema = "CREATE TABLE x (key TEXT PRIMARY KEY, value TEXT);".to_string();
+        // The hidden column is placed first to verify that column index handling
+        // remains correct when hidden columns are excluded from queries
+        // (e.g., in `*` expansion or `PRAGMA table_info`). It also includes a NOT NULL
+        // constraint and default value to confirm that SQLite silently ignores them
+        // on hidden columns.
+        let schema = "CREATE TABLE x (
+            comment TEXT HIDDEN NOT NULL DEFAULT 'default comment',
+            key TEXT PRIMARY KEY,
+            value TEXT
+        )"
+        .into();
         Ok((
             schema,
             KVStoreTable {
@@ -67,8 +77,9 @@ impl VTabCursor for KVStoreCursor {
                 log::debug!("idx_str found: key_eq\n value: {key:?}");
                 if let Some(key) = key {
                     let rowid = hash_key(&key);
-                    if let Some((k, v)) = self.store.borrow().get(&rowid) {
+                    if let Some((comment, k, v)) = self.store.borrow().get(&rowid) {
-                        self.rows.push((rowid, k.clone(), v.clone()));
+                        self.rows
+                            .push((rowid, comment.clone(), k.clone(), v.clone()));
                         self.index = Some(0);
                     } else {
                         self.rows.clear();
@@ -86,9 +97,9 @@ impl VTabCursor for KVStoreCursor {
                     .store
                     .borrow()
                     .iter()
-                    .map(|(&rowid, (k, v))| (rowid, k.clone(), v.clone()))
+                    .map(|(&rowid, (comment, k, v))| (rowid, comment.clone(), k.clone(), v.clone()))
                     .collect();
-                self.rows.sort_by_key(|(rowid, _, _)| *rowid);
+                self.rows.sort_by_key(|(rowid, _, _, _)| *rowid);
                 if self.rows.is_empty() {
                     self.index = None;
                     ResultCode::EOF
@@ -113,10 +124,11 @@ impl VTabCursor for KVStoreCursor {
         if self.index.is_some_and(|c| c >= self.rows.len()) {
             return Err("cursor out of range".into());
         }
-        if let Some((_, ref key, ref val)) = self.rows.get(self.index.unwrap_or(0)) {
+        if let Some((_, ref comment, ref key, ref val)) = self.rows.get(self.index.unwrap_or(0)) {
             match idx {
-                0 => Ok(Value::from_text(key.clone())), // key
+                0 => Ok(Value::from_text(comment.clone())),
-                1 => Ok(Value::from_text(val.clone())), // value
+                1 => Ok(Value::from_text(key.clone())), // key
+                2 => Ok(Value::from_text(val.clone())), // value
                 _ => Err("Invalid column".into()),
             }
         } else {
@@ -159,13 +171,13 @@ impl VTable for KVStoreTable {
         for constraint in constraints.iter() {
             if constraint.usable
                 && constraint.op == ConstraintOp::Eq
-                && constraint.column_index == 0
+                && constraint.column_index == 1
             {
                 // this extension wouldn't support order by but for testing purposes,
                 // we will consume it if we find an ASC order by clause on the value column
                 let mut consumed = false;
                 if let Some(order) = _order_by.first() {
-                    if order.column_index == 1 && !order.desc {
+                    if order.column_index == 2 && !order.desc {
                         consumed = true;
                     }
                 }
@@ -196,19 +208,24 @@ impl VTable for KVStoreTable {
     }
 
     fn insert(&mut self, values: &[Value]) -> Result<i64, Self::Error> {
-        let key = values
+        let comment = values
             .first()
             .and_then(|v| v.to_text())
+            .map(|v| v.to_string())
+            .unwrap_or("auto-generated".into());
+        let key = values
+            .get(1)
+            .and_then(|v| v.to_text())
             .ok_or("Missing key")?
             .to_string();
         let val = values
-            .get(1)
+            .get(2)
             .and_then(|v| v.to_text())
             .ok_or("Missing value")?
             .to_string();
         let rowid = hash_key(&key);
         {
-            self.store.borrow_mut().insert(rowid, (key, val));
+            self.store.borrow_mut().insert(rowid, (comment, key, val));
         }
         Ok(rowid)
     }

testing/cli_tests/extensions.py

@@ -797,6 +797,137 @@ def test_tablestats():
     limbo.quit()
 
 
+def test_hidden_columns():
+    _test_hidden_columns(exec_name=None, ext_path="target/debug/libturso_ext_tests")
+    _test_hidden_columns(exec_name="sqlite3", ext_path="target/debug/liblimbo_sqlite_test_ext")
+
+
+def _test_hidden_columns(exec_name, ext_path):
+    console.info(f"Running test_hidden_columns for {ext_path}")
+
+    limbo = TestTursoShell(exec_name=exec_name,)
+    limbo.execute_dot(f".load {ext_path}")
+    limbo.execute_dot(
+        "create virtual table t using kv_store;",
+    )
+    limbo.run_test_fn(".schema", lambda res: "CREATE VIRTUAL TABLE t" in res)
+    limbo.run_test_fn(
+        "insert into t(key, value) values ('k0', 'v0');",
+        null,
+        "can insert if hidden column is not specified explicitly",
+    )
+    limbo.run_test_fn(
+        "insert into t(key, value) values ('k1', 'v1');",
+        null,
+        "can insert if hidden column is not specified explicitly",
+    )
+    limbo.run_test_fn(
+        "select comment from t where key = 'k0';",
+        lambda res: "auto-generated" == res,
+        "can select a hidden column from kv_store",
+    )
+    limbo.run_test_fn(
+        "select comment from (select * from t where key = 'k0');",
+        lambda res: "Column comment not found" in res or "no such column: comment" in res,
+        "hidden columns are not exposed by subqueries by default",
+    )
+    limbo.run_test_fn(
+        "select * from (select comment from t where key = 'k0');",
+        lambda res: "auto-generated" == res,
+        "can select hidden column exposed by subquery",
+    )
+    limbo.run_test_fn(
+        "insert into t(comment, key, value) values ('my comment', 'hidden', 'test');",
+        null,
+        "can insert if a hidden column is specified explicitly",
+    )
+    limbo.run_test_fn(
+        "select comment from t where key = 'hidden';",
+        lambda res: "my comment" == res,
+        "can select a hidden column from kv_store",
+    )
+    limbo.run_test_fn(
+        "select * from t where key = 'hidden';",
+        lambda res: "hidden|test" == res,
+        "hidden column is excluded from * expansion",
+    )
+    limbo.run_test_fn(
+        "select t.* from t where key = 'hidden';",
+        lambda res: "hidden|test" == res,
+        "hidden column is excluded from <table name>.* expansion",
+    )
+    limbo.run_test_fn(
+        "insert into t(comment, key, value) values ('insert_hidden', 'test');",
+        lambda res: "2 values for 3 columns" in res,
+        "fails when number of values does not match number of specified columns",
+    )
+    limbo.run_test_fn(
+        "update t set comment = 'updated comment' where key = 'hidden';",
+        null,
+        "can update a hidden column if specified explicitly",
+    )
+    limbo.run_test_fn(
+        "select comment from t where key = 'hidden';",
+        lambda res: "updated comment" == res,
+    )
+    limbo.run_test_fn(
+        "PRAGMA table_info=t;",
+        lambda res: "0|key|TEXT|0|TURSO|1\n1|value|TEXT|0|TURSO|0" == res,
+        "hidden columns are not listed in the dataset returned by 'PRAGMA table_info'",
+    )
+    limbo.run_test_fn(
+        "select comment, count(*) from t group by comment;",
+        lambda res: "auto-generated|2\nupdated comment|1" == res,
+        "can use hidden columns in aggregations",
+    )
+
+    # ORDER BY
+    limbo.execute_dot("CREATE VIRTUAL TABLE o USING kv_store;")
+    limbo.run_test_fn(".schema", lambda res: "CREATE VIRTUAL TABLE o" in res)
+    limbo.execute_dot("INSERT INTO o(comment, key, value) VALUES ('0', '5', 'a');")
+    limbo.execute_dot("INSERT INTO o(comment, key, value) VALUES ('1', '4', 'b');")
+    limbo.execute_dot("INSERT INTO o(comment, key, value) VALUES ('2', '3', 'c');")
+    limbo.run_test_fn(
+        "SELECT * FROM o ORDER BY comment;",
+        lambda res: "5|a\n4|b\n3|c" == res,
+    )
+    limbo.run_test_fn(
+        "SELECT * FROM o ORDER BY 0;",
+        lambda res: "invalid column index: 0" in res or "term out of range - should be between 1 and 2" in res,
+    )
+    limbo.run_test_fn(
+        "SELECT * FROM o ORDER BY 1;",
+        lambda res: "3|c\n4|b\n5|a" == res,
+    )
+
+    # JOINs
+    limbo.execute_dot("CREATE TABLE r (comment, key, value);")
+    limbo.execute_dot("INSERT INTO r VALUES ('comment0', '2', '3');")
+    limbo.execute_dot("INSERT INTO r VALUES ('comment1', '4', '5');")
+    limbo.execute_dot("CREATE VIRTUAL TABLE l USING kv_store;")
+    limbo.run_test_fn(".schema", lambda res: "CREATE VIRTUAL TABLE l" in res)
+    limbo.execute_dot("INSERT INTO l(comment, key, value) values ('comment1', '2', '3');")
+    limbo.run_test_fn(
+        "SELECT * FROM l NATURAL JOIN r;",
+        lambda res: "2|3|comment0" == res,
+    )
+    limbo.run_test_fn(
+        "SELECT * FROM l JOIN r USING (comment);",
+        lambda res: "2|3|4|5" == res,
+    )
+    limbo.run_test_fn(
+        "SELECT * FROM l JOIN r ON l.comment = r.comment;",
+        lambda res: "2|3|comment1|4|5" == res,
+    )
+    # TODO: Limbo panics for:
+    # - SELECT * FROM l NATURAL JOIN r NATURAL JOIN r;
+    # - SELECT * FROM l NATURAL JOIN r NATURAL JOIN l;
+    # - SELECT * FROM r NATURAL JOIN l;
+    # - SELECT * FROM r NATURAL JOIN l NATURAL JOIN r;
+
+    limbo.quit()
+
+
 def main():
     try:
         test_regexp()
@@ -812,6 +943,7 @@ def main():
         test_create_virtual_table()
         test_csv()
         test_tablestats()
+        test_hidden_columns()
     except Exception as e:
         console.error(f"Test FAILED: {e}")
         cleanup()

testing/sqlite_test_ext/src/kvstore.c

@@ -10,6 +10,7 @@ SQLITE_EXTENSION_INIT1
 #include <string.h>
 
 typedef struct {
+    char *comment;
     char *key;
     char *value;
     sqlite3_int64 rowid;
@@ -38,7 +39,12 @@ static int kvstoreConnect(
     kv_table *pNew;
     int rc;
 
-    rc = sqlite3_declare_vtab(db, "CREATE TABLE x (key TEXT PRIMARY KEY, value TEXT)");
+    rc = sqlite3_declare_vtab(db,
+        "CREATE TABLE x("
+        "  comment TEXT HIDDEN NOT NULL DEFAULT 'default comment',"
+        "  key TEXT PRIMARY KEY,"
+        "  value TEXT"
+        ")");
 
     if (rc == SQLITE_OK) {
         pNew = sqlite3_malloc(sizeof(*pNew));
@@ -54,6 +60,7 @@ static int kvstoreConnect(
 static int kvstoreDisconnect(sqlite3_vtab *pVtab) {
     kv_table *table = (kv_table *)pVtab;
     for (int i = 0; i < table->row_count; i++) {
+        sqlite3_free(table->rows[i].comment);
         sqlite3_free(table->rows[i].key);
         sqlite3_free(table->rows[i].value);
     }
@@ -101,9 +108,12 @@ static int kvstoreColumn(sqlite3_vtab_cursor *cur, sqlite3_context *ctx, int col
     kv_row *row = &cursor->table->rows[cursor->current];
     switch (col) {
         case 0:
-            sqlite3_result_text(ctx, row->key, -1, SQLITE_TRANSIENT);
+            sqlite3_result_text(ctx, row->comment, -1, SQLITE_TRANSIENT);
             break;
         case 1:
+            sqlite3_result_text(ctx, row->key, -1, SQLITE_TRANSIENT);
+            break;
+        case 2:
             sqlite3_result_text(ctx, row->value, -1, SQLITE_TRANSIENT);
             break;
     }
@@ -140,13 +150,21 @@ static int kvUpsert(
 ) {
     kv_table *table = (kv_table *)pVTab;
 
-    const char *key = (const char *)sqlite3_value_text(argv[2]);
+    const char *comment;
-    const char *value = (const char *)sqlite3_value_text(argv[3]);
+    if (sqlite3_value_type(argv[2]) == SQLITE_NULL) {
+        comment = "auto-generated";
+    } else {
+        comment = (const char *)sqlite3_value_text(argv[2]);
+    }
+    const char *key = (const char *)sqlite3_value_text(argv[3]);
+    const char *value = (const char *)sqlite3_value_text(argv[4]);
 
     // Check if key exists; if so, replace
     for (int i = 0; i < table->row_count; i++) {
         if (strcmp(table->rows[i].key, key) == 0) {
+            sqlite3_free(table->rows[i].comment);
             sqlite3_free(table->rows[i].value);
+            table->rows[i].comment = sqlite3_mprintf("%s", comment);
             table->rows[i].value = sqlite3_mprintf("%s", value);
             return SQLITE_OK;
         }
@@ -155,6 +173,7 @@ static int kvUpsert(
     // Otherwise, insert new
     table->rows = sqlite3_realloc(table->rows, sizeof(kv_row) * (table->row_count + 1));
     kv_row *row = &table->rows[table->row_count++];
+    row->comment = sqlite3_mprintf("%s", comment);
     row->key = sqlite3_mprintf("%s", key);
     row->value = sqlite3_mprintf("%s", value);
     row->rowid = table->next_rowid;
@@ -175,6 +194,7 @@ static int kvDelete(sqlite3_vtab *pVTab, sqlite3_int64 rowid) {
     }
 
     if (idx > -1) {
+        sqlite3_free(table->rows[idx].comment);
         sqlite3_free(table->rows[idx].key);
         sqlite3_free(table->rows[idx].value);
 
@@ -196,7 +216,7 @@ static int kvstoreUpdate(
     if (argc == 1) {
         return kvDelete(pVTab, sqlite3_value_int64(argv[0]));
     } else {
-        assert(argc == 4);
+        assert(argc == 5);
         return kvUpsert(pVTab, argv, pRowid);
     }
 }