From cb275feaa29ff1c5f3a7cf8ef983bf935e99696a Mon Sep 17 00:00:00 2001
From: Ajaya Agrawal <ajaya@rippling.com>
Date: Sun, 1 Sep 2024 16:16:41 +0530
Subject: [PATCH] Index lookup support

Adds support for parsing index structure
---
 core/lib.rs    | 21 +++++++++++----
 core/schema.rs | 70 +++++++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 83 insertions(+), 8 deletions(-)

diff --git a/core/lib.rs b/core/lib.rs
index ee329b2f2..d4ea51fad 100644
--- a/core/lib.rs
+++ b/core/lib.rs
@@ -84,13 +84,24 @@ impl Database {
                 match rows.next_row()? {
                     RowResult::Row(row) => {
                         let ty = row.get::<&str>(0)?;
-                        if ty != "table" {
+                        if ty != "table" && ty != "index"{
                             continue;
                         }
-                        let root_page: i64 = row.get::<i64>(3)?;
-                        let sql: &str = row.get::<&str>(4)?;
-                        let table = schema::BTreeTable::from_sql(sql, root_page as usize)?;
-                        schema.add_table(Rc::new(table));
+                        match ty {
+                            "table" => {
+                                let root_page: i64 = row.get::<i64>(3)?;
+                                let sql: &str = row.get::<&str>(4)?;
+                                let table = schema::BTreeTable::from_sql(sql, root_page as usize)?;
+                                schema.add_table(Rc::new(table));
+                            }
+                            "index" => {
+                                let root_page: i64 = row.get::<i64>(3)?;
+                                let sql: &str = row.get::<&str>(4)?;
+                                let index = schema::Index::from_sql(sql, root_page as usize)?;
+                                schema.add_index(Rc::new(index));
+                            }
+                            _ => continue
+                        }
                     }
                     RowResult::IO => {
                         // TODO: How do we ensure that the I/O we submitted to
diff --git a/core/schema.rs b/core/schema.rs
index 851fab6f1..607d80c10 100644
--- a/core/schema.rs
+++ b/core/schema.rs
@@ -1,8 +1,9 @@
-use crate::{util::normalize_ident, Result};
+use crate::{util::normalize_ident, LimboError, Result};
 use core::fmt;
+use std::cell::RefCell;
 use fallible_iterator::FallibleIterator;
 use log::trace;
-use sqlite3_parser::ast::{Expr, Literal, TableOptions};
+use sqlite3_parser::ast::{Expr, IndexedColumn, Literal, TableOptions};
 use sqlite3_parser::{
     ast::{Cmd, CreateTableBody, QualifiedName, ResultColumn, Stmt},
     lexer::sql::Parser,
@@ -12,13 +13,16 @@ use std::rc::Rc;
 
 pub struct Schema {
     pub tables: HashMap<String, Rc<BTreeTable>>,
+    // table_name to list of indexes for the table
+    pub indexes: HashMap<String, Vec<Rc<Index>>>,
 }
 
 impl Schema {
     pub fn new() -> Self {
         let mut tables: HashMap<String, Rc<BTreeTable>> = HashMap::new();
+        let mut indexes: HashMap<String, Vec<Rc<Index>>> = HashMap::new();
         tables.insert("sqlite_schema".to_string(), Rc::new(sqlite_schema_table()));
-        Self { tables }
+        Self { tables, indexes}
     }
 
     pub fn add_table(&mut self, table: Rc<BTreeTable>) {
@@ -30,6 +34,14 @@ impl Schema {
         let name = normalize_ident(name);
         self.tables.get(&name).cloned()
     }
+
+    pub fn add_index(&mut self, index: Rc<Index>){
+        let table_name = normalize_ident(&index.table_name);
+        self.indexes
+            .entry(table_name)
+            .or_insert_with(Vec::new)
+            .push(index.clone())
+    }
 }
 
 #[derive(Clone, Debug)]
@@ -385,6 +397,58 @@ pub fn sqlite_schema_table() -> BTreeTable {
     }
 }
 
+#[derive(Debug)]
+pub struct Index {
+    pub name: String,
+    pub table_name: String,
+    pub root_page: usize,
+    pub columns: Vec<IndexColumn>,
+    pub unique: bool,
+}
+
+#[derive(Debug, Clone)]
+pub struct IndexColumn {
+    pub name: String,
+    pub order: Order,
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub enum Order {
+    Ascending,
+    Descending,
+}
+
+impl Index {
+    pub fn from_sql(sql: &str, root_page: usize) -> Result<Index>{
+        let mut parser = Parser::new(sql.as_bytes());
+        let cmd = parser.next()?;
+        match cmd {
+            Some(Cmd::Stmt(Stmt::CreateIndex {idx_name, tbl_name, columns, unique, ..})) => {
+                let index_name = normalize_ident(&idx_name.name.0);
+                let index_columns = columns.into_iter().map(|col| {
+                    IndexColumn {
+                        name: normalize_ident(&col.expr.to_string()),
+                        order: match col.order {
+                            Some(sqlite3_parser::ast::SortOrder::Asc) => Order::Ascending,
+                            Some(sqlite3_parser::ast::SortOrder::Desc) => Order::Descending,
+                            None => Order::Ascending,
+                        },
+                    }
+                }).collect();
+                Ok(Index {
+                    name: index_name,
+                    table_name: normalize_ident(&tbl_name.0),
+                    root_page: root_page,
+                    columns: index_columns,
+                    unique,
+                }
+                )
+            }
+            _ => todo!("Expected create index statement"),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;