Build your own SQLite

October 19, 2024

Bootstrapping

让我们从建立一个简单的是数据库文件开始（通过 sqlite3）

sqlite3 minimal_test.db
sqlite> create table table1(id integer);
sqlite> create table table2(id integer);
sqlite> .exit

这将创建两个表，只含 id 列。

sqlite3 minimal_test.db
sqlite> .tables
table1  table2
sqlite> .exit

我们创建一个 rust 项目，并添加 anyhow 为依赖：

cargo new rsqlite
cd rsqlite
cargo add anyhow

SQLite File Format

SQLite 的文件格式定义在 SQlite3 文档：Database File Format 中。

简单来说：

SQLite 数据库文件是单个文件，包含数据库的所有数据和元数据。
这个文件被组织成一系列的页 (page)，每个页的大小通常是 4KB (4096 字节)，但也可以是 512 字节到 64KB 之间的任何 2 次幂。

页 (Page) 类型

数据库头页 (Database Header Page): 数据库文件的第一页。包含了数据库的元数据，例如页大小、数据库格式版本、数据库 schema 等。
空闲页列表页 (Freelist Page): 这些页跟踪数据库中所有未使用的页。
B 树页 (B-tree Page): 这些页存储数据库中的实际数据，并使用 B 树为数据结构进行组织。
溢出页 (Overflow Page): 当 B 树页中的数据太大而无法容纳在一个页中时，可使用溢出页来存储额外的数据。

B 树页结构

B 树页是 SQLite 数据库文件中最常见的页类型，用于存储数据库中的表数据。每个 B 树页包含：

页头 (Page Header): 包含页码、页类型、指向父页和子页的指针等信息。
单元指针数组 (Cell Pointer Array): 指向页中存储的每个数据单元 (cell) 的指针数组。
单元内容区 (Cell Content Area): 存储实际数据单元 (cell) 的区域。

数据单元 (Cell)

每个数据单元 (cell) 代表数据库表中的一行数据，包含：

单元头 (Cell Header): 包含单元的大小、数据类型等信息。
数据 (Data): 实际的数据值。

这里我们将只对 table btree leaf 页面（B 树页）感兴趣，因为它存储的是实际的表数据。

实现

SQLite 数据库存储在单个文件中，文件分为多页，每页大小相同：2 的幂次，介于 512 和 65536 字节之间。
第一页的前 100 个字节包含数据库头，其中包括页面大小和文件格式版本等信息。

我们的首要任务是实现一个 Pager 结构，用于读取和缓存数据库文件中的页面。
但在此之前，我们必须从数据库头中读取页面大小。

我们定义这些需要的结构如下：

page.rs

// 数据库头，可读取页面大小
#[derive(Debug, Copy, Clone)]
pub struct DbHeader {
    pub page_size: u32,
}

// 页面可以有多种类型，这里我们只对 table btree leaf 页面感兴趣，因为它存储的是实际的表数据。
#[derive(Debug, Copy, Clone)]
pub enum PageType {
    TableLeaf,
}

#[derive(Debug, Copy, Clone)]
pub struct PageHeader {
    // 表示页面类型的字节。对于 table btree-leaf 页面，它是 0x0D。
    pub page_type: PageType,
    // 2 字节整数，代表页面中第一个空闲区块的偏移量，如果没有空闲区块，则为 0。
    pub first_freeblock: u16,
    // 2 字节整数，代表页面中的单元格数。
    pub cell_count: u16,
    // 2 字节整数，代表第一个单元格的偏移量。
    pub cell_content_offset: u32,
    // 1 字节整数，代表页面中零散的空闲字节数（我们暂时不会使用它）。
    pub fragmented_bytes_count: u8,
}

#[derive(Debug, Clone)]
pub enum Page {
    TableLeaf(TableLeafPage),
}

// table btree-leaf 页以一个 8 字节的 header 开始，
// 随后是一串 "单元格指针"，其中包含页面中每个单元格的偏移量。
// 单元格包含表数据，我们可以将其视为键值对，
// 其中键是一个以 varint (rowid) （即可变长度）编码的 64 位整数，
// 值是代表行数据的任意字节序列。
#[derive(Debug, Clone)]
pub struct TableLeafPage {
    pub header: PageHeader,
    pub cell_pointers: Vec<u16>,
    pub cells: Vec<TableLeafCell>,
}

// 每个单元格包含表中一行的值，使用 SQLite 记录格式编码。
#[derive(Debug, Clone)]
pub struct TableLeafCell {
    pub size: i64,
    pub row_id: i64,
    pub payload: Vec<u8>,
}

我们实现一系列解析函数，最终组装出目标结构。

pager.rs

use std::{
    collections::HashMap,
    io::{Read, Seek, SeekFrom},
};

use anyhow::Context;

use crate::page;

pub const HEADER_SIZE: usize = 100;
const HEADER_PREFIX: &[u8] = b"SQLite format 3\0";
const HEADER_PAGE_SIZE_OFFSET: usize = 16;

const PAGE_MAX_SIZE: u32 = 65536;
const PAGE_LEAF_HEADER_SIZE: usize = 8;

const PAGE_LEAF_TABLE_ID: u8 = 13;

const PAGE_FIRST_FREEBLOCK_OFFSET: usize = 1;
const PAGE_CELL_COUNT_OFFSET: usize = 3;
const PAGE_CELL_CONTENT_OFFSET: usize = 5;
const PAGE_FRAGMENTED_BYTES_COUNT_OFFSET: usize = 7;

pub fn parse_header(buffer: &[u8]) -> anyhow::Result<page::DbHeader> {
    // 数据库头以 magic string `SQLite format 3\0` 开头
    if !buffer.starts_with(HEADER_PREFIX) {
        let prefix = String::from_utf8_lossy(&buffer[..HEADER_PREFIX.len()]);
        anyhow::bail!("invalid header prefix: {prefix}");
    }

    // 紧跟着一个 2 字节大端整数，表示页面大小（是2的幂次）
    let page_size_raw = read_be_word_at(buffer, HEADER_PAGE_SIZE_OFFSET);
    let page_size = match page_size_raw {
        1 => PAGE_MAX_SIZE, // 由于最大页面大小无法用 2 字节整数表示，因此使用页面大小为 1 来表示最大页面大小
        n if n.is_power_of_two() => n as u32,
        _ => anyhow::bail!("page size is not a power of 2: {}", page_size_raw),
    };

    Ok(page::DbHeader { page_size })
}

fn parse_page(buffer: &[u8], page_num: usize) -> anyhow::Result<page::Page> {
    // 第一页包含数据库头，此时从偏移量 100 开始解析页面
    let ptr_offset = if page_num == 1 { HEADER_SIZE as u16 } else { 0 };

    match buffer[0] {
        PAGE_LEAF_TABLE_ID => parse_table_leaf_page(buffer, ptr_offset),
        _ => Err(anyhow::anyhow!("unknown page type: {}", buffer[0])),
    }
}

fn parse_table_leaf_page(buffer: &[u8], ptr_offset: u16) -> anyhow::Result<page::Page> {
    let header = parse_page_header(buffer)?;

    let content_buffer = &buffer[PAGE_LEAF_HEADER_SIZE..];
    let cell_pointers = parse_cell_pointers(content_buffer, header.cell_count as usize, ptr_offset);

    let cells = cell_pointers
        .iter()
        .map(|&ptr| parse_table_leaf_cell(&buffer[ptr as usize..]))
        .collect::<anyhow::Result<Vec<page::TableLeafCell>>>()?;

    Ok(page::Page::TableLeaf(page::TableLeafPage {
        header,
        cell_pointers,
        cells,
    }))
}

fn parse_table_leaf_cell(mut buffer: &[u8]) -> anyhow::Result<page::TableLeafCell> {
    let (n, size) = read_varint_at(buffer, 0);
    buffer = &buffer[n as usize..];

    let (n, row_id) = read_varint_at(buffer, 0);
    buffer = &buffer[n as usize..];

    let payload = buffer[..size as usize].to_vec();

    Ok(page::TableLeafCell {
        size,
        row_id,
        payload,
    })
}

fn parse_page_header(buffer: &[u8]) -> anyhow::Result<page::PageHeader> {
    let page_type = match buffer[0] {
        0x0d => page::PageType::TableLeaf,
        _ => anyhow::bail!("unknown page type: {}", buffer[0]),
    };

    let first_freeblock = read_be_word_at(buffer, PAGE_FIRST_FREEBLOCK_OFFSET);
    let cell_count = read_be_word_at(buffer, PAGE_CELL_COUNT_OFFSET);
    let cell_content_offset = match read_be_word_at(buffer, PAGE_CELL_CONTENT_OFFSET) {
        0 => 65536,
        n => n as u32,
    };
    let fragmented_bytes_count = buffer[PAGE_FRAGMENTED_BYTES_COUNT_OFFSET];

    Ok(page::PageHeader {
        page_type,
        first_freeblock,
        cell_count,
        cell_content_offset,
        fragmented_bytes_count,
    })
}

fn parse_cell_pointers(buffer: &[u8], n: usize, ptr_offset: u16) -> Vec<u16> {
    let mut pointers = Vec::with_capacity(n);
    for i in 0..n {
        pointers.push(read_be_word_at(buffer, 2 * i) - ptr_offset);
    }
    pointers
}

// Varint 的编码规则是：
// 每个字节的最高位 (MSB) 用作标志位：
// 1 表示还有后续字节
// 0 表示这是最后一个字节
// 每个字节的低 7 位存储整数的一部分，从低位到高位排列
pub fn read_varint_at(buffer: &[u8], mut offset: usize) -> (u8, i64) {
    let mut size = 0;
    let mut result = 0;

    // varint 最多占用 8 个字节。当前字节的最高位为 0，表示这是最后一个字节
    while size < 8 && buffer[offset] >= 0b1000_0000 {
        result |= ((buffer[offset] as i64) & 0b0111_1111) << (7 * size);
        offset += 1;
        size += 1;
    }

    result |= (buffer[offset] as i64) << (7 * size);

    (size + 1, result)
}

fn read_be_double_at(input: &[u8], offset: usize) -> u32 {
    u32::from_be_bytes(input[offset..offset + 4].try_into().unwrap())
}

fn read_be_word_at(input: &[u8], offset: usize) -> u16 {
    u16::from_be_bytes(input[offset..offset + 2].try_into().unwrap())
}

#[derive(Debug, Clone)]
pub struct Pager<I: Read + Seek = std::fs::File> {
    input: I,
    page_size: usize,
    pages: HashMap<usize, page::Page>,
}

impl<I: Read + Seek> Pager<I> {
    pub fn new(input: I, page_size: usize) -> Self {
        Self {
            input,
            page_size,
            pages: HashMap::new(),
        }
    }

    pub fn read_page(&mut self, n: usize) -> anyhow::Result<&page::Page> {
        if self.pages.contains_key(&n) {
            return Ok(self.pages.get(&n).unwrap());
        }

        let page = self.load_page(n)?;
        self.pages.insert(n, page);
        Ok(self.pages.get(&n).unwrap())
    }

    fn load_page(&mut self, n: usize) -> anyhow::Result<page::Page> {
        let offset = HEADER_SIZE + n.saturating_sub(1) * self.page_size;

        self.input
            .seek(SeekFrom::Start(offset as u64))
            .context("seek to page start")?;

        let mut buffer = vec![0; self.page_size];
        self.input.read_exact(&mut buffer).context("read page")?;

        parse_page(&buffer, n)
    }
}

我们现在有了读取页面和访问页面单元格的方法。但如何解码单元格的值呢？

每个单元格包含表中一行的值，使用 SQLite 记录（Record）格式编码。
记录格式非常简单：记录由标题和字段值序列组成。
header 以一个代表其字节大小的 varint 开始，后面是一个 varint 序列（每列一个）用于表示类型。

cursor.rs

use std::borrow::Cow;

use crate::{page::Page, pager::Pager, value::Value};

#[derive(Debug, Copy, Clone)]
pub enum RecordFieldType {
    Null,
    I8,
    I16,
    I24,
    I32,
    I48,
    I64,
    Float,
    Zero,
    One,
    String(usize),
    Blob(usize),
}

#[derive(Debug, Clone)]
pub struct RecordField {
    pub offset: usize,
    pub field_type: RecordFieldType,
}

#[derive(Debug, Clone)]
pub struct RecordHeader {
    pub fields: Vec<RecordField>,
}

fn parse_record_header(mut buffer: &[u8]) -> anyhow::Result<RecordHeader> {
    // header 以一个代表其字节大小的 varint 开始，后面是一个 varint 序列
    let (varint_size, header_length) = crate::pager::read_varint_at(buffer, 0);
    buffer = &buffer[varint_size as usize..header_length as usize];

    let mut fields = Vec::new();
    let mut current_offset = header_length as usize;

    while !buffer.is_empty() {
        let (discriminant_size, discriminant) = crate::pager::read_varint_at(buffer, 0);
        buffer = &buffer[discriminant_size as usize..];

        // 确定每一列的类型
        let (field_type, field_size) = match discriminant {
            0 => (RecordFieldType::Null, 0),
            1 => (RecordFieldType::I8, 1),
            2 => (RecordFieldType::I16, 2),
            3 => (RecordFieldType::I24, 3),
            4 => (RecordFieldType::I32, 4),
            5 => (RecordFieldType::I48, 6),
            6 => (RecordFieldType::I64, 8),
            7 => (RecordFieldType::Float, 8), // 64 位 IEEE 浮点数
            8 => (RecordFieldType::Zero, 0),  // 整数 0
            9 => (RecordFieldType::One, 0),   // 整数 1
            // 10 和 11：保留给内部使用
            n if n >= 12 && n % 2 == 0 => {
                // n 偶数且 n > 12：大小为 (n - 12) / 2 的 BLOB
                let size = ((n - 12) / 2) as usize;
                (RecordFieldType::Blob(size), size)
            }
            n if n >= 13 && n % 2 == 1 => {
                // n 为奇数且 n > 13：大小为 (n - 13) / 2 的 TEXT
                let size = ((n - 13) / 2) as usize;
                (RecordFieldType::String(size), size)
            }
            n => anyhow::bail!("unsupported field type: {}", n),
        };

        fields.push(RecordField {
            offset: current_offset,
            field_type,
        });

        current_offset += field_size;
    }

    Ok(RecordHeader { fields })
}

为了更方便地处理记录，我们将定义一个 Value 类型（表示字段值）和一个 Cursor 结构（唯一标识数据库文件中的记录）。
Cursor 将暴露一个 field 方法，返回记录第 n 个字段的值：

cursor.rs

#[derive(Debug)]
pub struct Cursor<'p> {
    header: RecordHeader,
    pager: &'p mut Pager,
    page_index: usize,
    page_cell: usize,
}

impl<'p> Cursor<'p> {
    pub fn field(&mut self, n: usize) -> Option<Value> {
        let record_field = self.header.fields.get(n)?;

        let payload = match self.pager.read_page(self.page_index) {
            Ok(Page::TableLeaf(leaf)) => &leaf.cells[self.page_cell].payload,
            _ => return None,
        };

        match record_field.field_type {
            RecordFieldType::Null => Some(Value::Null),
            RecordFieldType::I8 => Some(Value::Int(read_i8_at(payload, record_field.offset))),
            RecordFieldType::I16 => Some(Value::Int(read_i16_at(payload, record_field.offset))),
            RecordFieldType::I24 => Some(Value::Int(read_i24_at(payload, record_field.offset))),
            RecordFieldType::I32 => Some(Value::Int(read_i32_at(payload, record_field.offset))),
            RecordFieldType::I48 => Some(Value::Int(read_i48_at(payload, record_field.offset))),
            RecordFieldType::I64 => Some(Value::Int(read_i64_at(payload, record_field.offset))),
            RecordFieldType::Float => Some(Value::Float(read_f64_at(payload, record_field.offset))),
            RecordFieldType::String(length) => {
                let value = std::str::from_utf8(
                    &payload[record_field.offset..record_field.offset + length],
                )
                .expect("invalid utf8");
                Some(Value::String(Cow::Borrowed(value)))
            }
            RecordFieldType::Blob(length) => {
                let value = &payload[record_field.offset..record_field.offset + length];
                Some(Value::Blob(Cow::Borrowed(value)))
            }
            _ => panic!("unimplemented"),
        }
    }
}

fn read_i8_at(input: &[u8], offset: usize) -> i64 {
    input[offset] as i64
}

fn read_i16_at(input: &[u8], offset: usize) -> i64 {
    i16::from_be_bytes(input[offset..offset + 2].try_into().unwrap()) as i64
}

fn read_i24_at(input: &[u8], offset: usize) -> i64 {
    (i32::from_be_bytes(input[offset..offset + 3].try_into().unwrap()) & 0x00FFFFFF) as i64
}

fn read_i32_at(input: &[u8], offset: usize) -> i64 {
    i32::from_be_bytes(input[offset..offset + 4].try_into().unwrap()) as i64
}

fn read_i48_at(input: &[u8], offset: usize) -> i64 {
    i64::from_be_bytes(input[offset..offset + 6].try_into().unwrap()) & 0x0000FFFFFFFFFFFF
}

fn read_i64_at(input: &[u8], offset: usize) -> i64 {
    i64::from_be_bytes(input[offset..offset + 8].try_into().unwrap())
}

fn read_f64_at(input: &[u8], offset: usize) -> f64 {
    f64::from_be_bytes(input[offset..offset + 8].try_into().unwrap())
}

value.rs

use std::borrow::Cow;

#[derive(Debug, Clone)]
pub enum Value<'p> {
    Null,
    String(Cow<'p, str>),
    Blob(Cow<'p, [u8]>),
    Int(i64),
    Float(f64),
}

impl<'p> Value<'p> {
    pub fn as_str(&self) -> Option<&str> {
        if let Value::String(s) = self {
            Some(s.as_ref())
        } else {
            None
        }
    }
}

为了简化对页面记录的迭代，我们还将实现一个 Scanner 结构，它可以封装页面，并允许我们为每条记录获取一个 Cursor 结构：

cursor.rs


#[derive(Debug)]
pub struct Scanner<'p> {
    pager: &'p mut Pager,
    page: usize,
    cell: usize,
}

impl<'p> Scanner<'p> {
    pub fn new(pager: &'p mut Pager, page: usize) -> Scanner<'p> {
        Scanner {
            pager,
            page,
            cell: 0,
        }
    }
    pub fn next_record(&mut self) -> Option<anyhow::Result<Cursor>> {
        let page = match self.pager.read_page(self.page) {
            Ok(page) => page,
            Err(e) => return Some(Err(e)),
        };

        match page {
            Page::TableLeaf(leaf) => {
                let cell = leaf.cells.get(self.cell)?;

                let header = match parse_record_header(&cell.payload) {
                    Ok(header) => header,
                    Err(e) => return Some(Err(e)),
                };

                let record = Cursor {
                    header,
                    pager: self.pager,
                    page_index: self.page,
                    page_cell: self.cell,
                };

                self.cell += 1;

                Some(Ok(record))
            }
        }
    }
}

在完成了大部分前期工作后，我们就可以回到最初的目标：列出表。
SQLite 将数据库的模式（Schema）存储在一个名为 sqlite_schema 的特殊表中。
该表用于表示 schema 对象。

sql

CREATE TABLE sqlite_schema(
  type text,
  name text,
  tbl_name text,
  rootpage integer,
  sql text
);

表字段的含义如下：

type：可为 table，index，view 或 trigger。我们这里只有表，因此只会是 table
name：schema 对象的名称
tbl_name：schema 对象关联的表名称，如果是表，则与 name 相同
rootpage：表的根页面，稍后我们将使用它来读取表的内容
sql：用于创建表的 SQL 语句

sqlite> select * from sqlite_schema;
table|table1|table1|2|CREATE TABLE table1(id integer)
table|table2|table2|3|CREATE TABLE table2(id integer)

由于我们的简单数据库目前只处理基本 schema，因此我们可以假定整个 schema 都包含在数据库文件的第一页中。
为了列出数据库中的表，我们需要：

用数据库文件初始化 Pager
为第一页创建 Scanner
遍历 records，并打印其 name 字段的值（位于索引 1）

db.rs

use std::{io::Read, path::Path};

use anyhow::Context;

use crate::{cursor::Scanner, page::DbHeader, pager, pager::Pager};

pub struct Db {
    pub header: DbHeader,
    pager: Pager,
}

impl Db {
    pub fn from_file(filename: impl AsRef<Path>) -> anyhow::Result<Db> {
        let mut file = std::fs::File::open(filename.as_ref()).context("open db file")?;

        let mut header_buffer = [0; pager::HEADER_SIZE];
        file.read_exact(&mut header_buffer)
            .context("read db header")?;

        let header = pager::parse_header(&header_buffer).context("parse db header")?;

        let pager = Pager::new(file, header.page_size as usize);

        Ok(Db { header, pager })
    }

    pub fn scanner(&mut self, page: usize) -> Scanner {
        Scanner::new(&mut self.pager, page)
    }
}

实现支持 tables 命令的基本 REPL 非常简单：

main.rs

use std::io::{stdin, BufRead, Write};

use anyhow::Context;

mod cursor;
mod db;
mod page;
mod pager;
mod value;

fn main() -> anyhow::Result<()> {
    let database = db::Db::from_file(std::env::args().nth(1).context("missing db file")?)?;
    cli(database)
}

fn cli(mut db: db::Db) -> anyhow::Result<()> {
    print_flushed("rqlite> ")?;

    let mut line_buffer = String::new();

    while stdin().lock().read_line(&mut line_buffer).is_ok() {
        match line_buffer.trim() {
            ".exit" => break,
            ".tables" => display_tables(&mut db)?,
            _ => {
                println!("Unrecognized command '{}'", line_buffer.trim());
            }
        }

        print_flushed("\nrqlite> ")?;

        line_buffer.clear();
    }

    Ok(())
}

fn display_tables(db: &mut db::Db) -> anyhow::Result<()> {
    let mut scanner = db.scanner(1);

    while let Some(Ok(mut record)) = scanner.next_record() {
        let type_value = record
            .field(0)
            .context("missing type field")
            .context("invalid type field")?;

        if type_value.as_str() == Some("table") {
            let name_value = record
                .field(1)
                .context("missing name field")
                .context("invalid name field")?;

            print!("{} ", name_value.as_str().unwrap());
        }
    }

    Ok(())
}

fn print_flushed(s: &str) -> anyhow::Result<()> {
    print!("{}", s);
    std::io::stdout().flush().context("flush stdout")
}