Go 語言 CGO 用戶深度定製 SQLite 代碼

本文是 BRUNO CALZA 記錄的關於如何改變 SQLite 源代碼，使記錄行更新時可用於 Go 的更新鉤子函數的過程。原文通過深度定製 C 語言的 API 函數達成目的，這幾乎是所有 CGO 深度用戶必然經歷的過程（關於 CGO 的基本用法可以參考譯者的《Go 高級編程》第 2 章），是一個非常有借鑑意義的技術文章。

作者：BRUNO CALZA
譯者：柴樹杉，凹語言作者、Go 語言貢獻者、多本 Go 語言圖書作者，目前在螞蟻從事 KusionStack 和 KCL 開發。
原文：https://brunocalza.me/making-a-change-to-sqlite-source-code/

背景

有一天，我正在考慮如何在 SQLite 中獲取最近插入或更新的行記錄的數據。這樣做的動機是我想創建該行的 hash，本質上是爲了在插入或更新行時能夠構建相應表的 Merkle 樹。

SQLite 提供的最符合的 API 可能是 sqlite3_update_hook:

sqlite3_update_hook() 函數爲數據庫連接註冊一個回調函數，該數據庫連接由第一個參數標識，在 rowwid 表中更新、插入或刪除行時調用。

這個 API 的問題是它只返回行的 rowid。這意味着還需要爲列內的行獲取所有列。即使使用這種方法，我仍然無法獲得行記錄的原始數據。只能得到那一行的驅動信息。

關於如何構建這樣的樹可能有很多方法，但就我而言 SQLite API 並沒有提供真正想要的東西。因此，我決定趁此機會更深入地挖掘下源代碼，同時看看內部實現的細節。不僅如此，我希望可以對它進行一些修改和測試，看看能否滿足需求。

因爲對 C 語言的畏懼，開始我只是想假裝看下幾個源文件就跑路。沒想到這次真的有驚喜。

看看 SQLite 的代碼結構

首先使用 fossil 工具克隆了 SQLite 源代碼，下面是文件。

如果你對數據庫比較熟悉，或許可以猜測出一些文件對應的操作。因此，我決定直接跳到 insert.c 文件，看看能不能找到一些有趣的東西。

遍歷函數名列表，路過 sqlite3Insert 函數，看到以下注釋：

** This routine is called to handle SQL of the following forms:
**
**    insert into TABLE (IDLIST) values(EXPRLIST),(EXPRLIST),...
**    insert into TABLE (IDLIST) select
**    insert into TABLE (IDLIST) default values
**

也許在這個函數中有一些可鼓搗的地方。我能夠對其中發生的情況進行一些猜測，但引起我注意的是對名稱類似於 sqlite3vdbeXXX 的函數的函數調用的數量。

這讓我想起 SQLite 底層使用了一個名爲 vdbe 的虛擬機。這意味着所有 SQL 語句都首先被翻譯成該虛擬機的語言。然後，執行引擎執行虛擬機代碼。讓我們看一個簡單的 INSERT 語句如何被翻譯成字節碼:

sqlite> create table a (a int, b text);
sqlite> explain INSERT INTO a VALUES (1, 'Hello');
addr  opcode         p1    p2    p3    p4             p5  comment      
----  -------------  ----  ----  ----  -------------  --  -------------
0     Init           0     8     0                    0   Start at 8
1     OpenWrite      0     2     0     2              0   root=2 iDb=0; a
2     Integer        1     2     0                    0   r[2]=1
3     String8        0     3     0     Hello          0   r[3]='Hello'
4     NewRowid       0     1     0                    0   r[1]=rowid
5     MakeRecord     2     2     4     DB             0   r[4]=mkrec(r[2..3])
6     Insert         0     4     1     a              57  intkey=r[1] data=r[4]
7     Halt           0     0     0                    0   
8     Transaction    0     1     1     0              1   usesStmtJournal=0
9     Goto           0     1     0                    0

我得出的結論是 sqlite3Insert 實際上是根據 SQLite 插入規則，將解析後的 INSERT 語句轉換爲一系列虛擬機字節碼指令。

因此這並不是我要找的地方。我真正需要的是在插入之前創建記錄的位置。我猜測那隻能是執行虛擬機代碼的地方，可能是執行 Insert (OP_INSERT) 操作碼的地方。

根據上圖我直接找到了 vdbe.c 文件的位置，直奔主題。

我發現有一個有 8000 行代碼的 switch( pOp->opcode ) 語句，通過 OP_INSERT 關鍵字找到插入操作對應的代碼位置。

在對應分支的第一行中，總算找到了相關的線索:

Mem *pData;       /* MEM cell holding data for the record to be inserted */

所以 pData 指向要插入的記錄數據。您可以在 L5402 中看到pData = &aMem[pOp->p2];，它是如何將 pData 值設置爲虛擬機內存 aMem 地址的，該地址位於虛擬機寄存器 p2 所指向的位置。

快速回顧一下: 首先在 insert.c 文件我們瞭解到 INSERT 語句被翻譯成一堆虛擬機指令。然後通過 INSERT 的數據通過這些sqlite3vdbeXXX 調用到達虛擬機。我假設將 OP_INSERT 操作碼和數據註冊到虛擬機是在第 2593 行:

sqlite3VdbeAddOp3(v, OP_Insert, iDataCur, aRegIdx[i], regNewData);

下面 regNewData 的一個更詳細的說明:

** The regNewData parameter is the first register in a range that contains
** the data to be inserted or the data after the update.  There will be
** pTab->nCol+1 registers in this range.  The first register (the one
** that regNewData points to) will contain the new rowid, or NULL in the
** case of a WITHOUT ROWID table.  The second register in the range will
** contain the content of the first table column.  The third register will
** contain the content of the second table column.  And so forth.
**
** The regOldData parameter is similar to regNewData except that it contains
** the data prior to an UPDATE rather than afterwards.  regOldData is zero
** for an INSERT.  This routine can distinguish between UPDATE and INSERT by
** checking regOldData for zero.

所以，在這一點上，我們正在用數據執行機器代碼。代碼向下滾動一點，讓我們看看如何使用 pData。在 L5448-L5449 處可以看到:

x.pData = pData->z;
x.nData = pData->n;

x 的定義如下：

BtreePayload x;   /* Payload to be inserted */

完美。再向下滾動一點，我們看到:

rc = sqlite3BtreeInsert(pC->uc.pCursor, &x,
(pOp->p5 & (OPFLAG_APPEND|OPFLAG_SAVEPOSITION|OPFLAG_PREFORMAT)),
seekResult
);

我們終於找到了插入原始數據的位置。但是，我們怎麼知道它的格式和這裏記錄的一樣呢? 如果仔細查看示例 INSERT 中的虛擬機代碼，在INSERT 操作碼之前有一個 MakeRecord 操作碼，它負責構建記錄。

你可以在 vdb.c 文件中查看 OP_MakeRecord 實現，並看到以下注釋:

You can check the OP_MakeRecord implementation at vdbe.c file and see the following comment:

將 P1 開頭的 P2 寄存器轉換爲記錄格式，用作數據庫表中的數據記錄或索引中的鍵。

在 case 語句的最後幾行看到了關鍵部分:

/* Invoke the update-hook if required. */
if( rc ) goto abort_due_to_error;
if( pTab ){
assert( db->xUpdateCallback!=0 );
assert( pTab->aCol!=0 );
db->xUpdateCallback(db->pUpdateArg,
(pOp->p5 & OPFLAG_ISUPDATE) ? SQLITE_UPDATE : SQLITE_INSERT,
zDb, pTab->zName, x.nKey);
}
break;

看來我需要的東西都在這裏了。更新鉤子鉤子和原始數據。只需要更新時傳遞給回調函數即可。

開始定製 SQLite

這就是我期望的 API：

db->xUpdateCallback(db->pUpdateArg,
	(pOp->p5 & OPFLAG_ISUPDATE) ? SQLITE_UPDATE : SQLITE_INSERT,
	zDb, pTab->zName, x.nKey, pData->z, pData->n);

傳遞的是數據（pData->z）和其大小 (pData->n)。

爲了解釋函數簽名的變化，還需要在多個地方進行相應的修改。

以下是 fossil 工具提示的變化的源文件：

EDITED     src/main.c
EDITED     src/sqlite.h.in
EDITED     src/sqlite3ext.h
EDITED     src/sqliteInt.h
EDITED     src/tclsqlite.c
EDITED     src/vdbe.c

還有一些針對編譯提示的修改。

克隆一份 Go SQLite 驅動

現在是時候在一個 Go 程序中創建一個簡單的測試了。我比較熟悉與 SQLite 交互的 mattn/go-sqlite3 驅動程序。該項目通過導入 SQLite 合併文件並通過 CGO 綁定工作。

因此還需要再克隆下 Go SQLite 驅動，更新被我修改的文件。並在 Go API 中進行了必要的更新以訪問新值。

主要是對 updateHookTrampoline 的更改，現在接收記錄爲 *C.Char 和 int 類型的數據大小，轉型爲字節 Slice 並將其傳遞給回調函數:

func updateHookTrampoline(handle unsafe.Pointer, op int, db *C.char, table *C.char, rowid int64, data *C.char, size int) {
	callback := lookupHandle(handle).(func(int, string, string, int64, []byte))
	callback(op, C.GoString(db), C.GoString(table), rowid, C.GoBytes(unsafe.Pointer(data), C.int(size)))
}

RegisterUpdateHook 函數也需要做同樣的調整。

改動後的效果

現在已經準備好了測試的所有東西。讓我們運行一個簡單的例子，靈感來自 SQLite Internals: Pages & B-trees 博客文章。

package main
import (
	"database/sql"
	"fmt"
	"log"
	"os"
	"github.com/mattn/go-sqlite3"
)
func main() {
	sqlite3conn := []*sqlite3.SQLiteConn{}
	sql.Register("sqlite3_with_hook_example",
		&sqlite3.SQLiteDriver{
			ConnectHook: func(conn *sqlite3.SQLiteConn) error {
				sqlite3conn = append(sqlite3conn, conn)
				conn.RegisterUpdateHook(func(op int, db string, table string, rowid int64, data []byte) {
					switch op {
					case sqlite3.SQLITE_INSERT:
						fmt.Printf("%x\n", data)
					}
				})
				return nil
			},
		})
	os.Remove("./foo.db")
	srcDb, err := sql.Open("sqlite3_with_hook_example", "./foo.db")
	if err != nil {
		log.Fatal(err)
	}
	defer srcDb.Close()
	srcDb.Ping()
	_, err = srcDb.Exec(`CREATE TABLE sandwiches (
		id INTEGER PRIMARY KEY,
		name TEXT,
		length REAL,
		count INTEGER
	);`)
	if err != nil {
		log.Fatal(err)
	}
	_, err = srcDb.Exec("INSERT INTO sandwiches (name, length, count) VALUES ('Italian', 7.5, 2);")
	if err != nil {
		log.Fatal(err)
	}
}

不要忘記添加更新 go.mod 文件 replace github.com/mattn/go-sqlite3 => github.com/brunocalza/go-sqlite3 v0.0.0-20220926005737-36475033d841，重新定向驅動。

運行後應該得到以下的結果：

05001b07014974616c69616e401e00000000000002

這正是 ('Italian', 7.5, 2) 數據的 Efficient Sandwich 編碼的結果，不包含主鍵和記錄的長度 (前兩個字節)。

看到輸出結果我才發現能夠理解 SQLite 源代碼的部分內容真的很有趣，儘管我不理解它的大部分。但是我做了一些更改並看到這些更改，並通過 Go 的驅動程序看到結果的變化。

老實說這種更改數據庫源代碼的方法風險太大。與新版本保持同步也是一個太大的問題，但這是一個值得記錄的有趣經歷。

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/BJvSQE7QziiAbSqtZjHkpw

猜你喜歡