Integrate with pg_mooncake

Author: dpxcc

.github/workflows/MainDistributionPipeline.yml

@@ -29,6 +29,7 @@ jobs:
       duckdb_version: v1.3.2
       ci_tools_version: v1.3.2
       enable_rust: true
+      exclude_archs: windows_amd64;windows_amd64_mingw;wasm_mvp;wasm_eh;wasm_threads
       extension_name: mooncake
 
   code-quality-check:
@@ -38,4 +39,4 @@ jobs:
       duckdb_version: v1.3.2
       ci_tools_version: main
       extension_name: mooncake
-      format_checks: 'format;tidy'
+      format_checks: format

src/include/moonlink/moonlink.hpp

@@ -11,7 +11,7 @@ class Moonlink {
 
 	DataPtr GetTableSchema(const string &schema, const string &table);
 
-	DataPtr ScanTableBegin(const string &schema, const string &table);
+	DataPtr ScanTableBegin(const string &schema, const string &table, uint64_t lsn);
 
 	void ScanTableEnd(const string &schema, const string &table);
 

src/include/pgmooncake.hpp

@@ -0,0 +1,36 @@
+#pragma once
+
+#include "duckdb/common/dl.hpp"
+
+namespace duckdb {
+
+class Pgmooncake {
+public:
+	using drop_cstring_fn = void (*)(char *);
+	using get_init_query_fn = char *(*)();
+	using get_lsn_fn = uint64_t (*)();
+
+	static string GetInitQuery() {
+		static drop_cstring_fn drop_cstring =
+		    reinterpret_cast<drop_cstring_fn>(dlsym(RTLD_DEFAULT, "pgmooncake_drop_cstring"));
+		static get_init_query_fn get_init_query =
+		    reinterpret_cast<get_init_query_fn>(dlsym(RTLD_DEFAULT, "pgmooncake_get_init_query"));
+
+		if (!get_init_query) {
+			return "";
+		}
+		char *init_query = get_init_query();
+		string res(init_query);
+		D_ASSERT(drop_cstring);
+		drop_cstring(init_query);
+		return res;
+	}
+
+	static uint64_t GetLsn() {
+		static get_lsn_fn get_lsn = reinterpret_cast<get_lsn_fn>(dlsym(RTLD_DEFAULT, "pgmooncake_get_lsn"));
+
+		return get_lsn ? get_lsn() : 0;
+	}
+};
+
+} // namespace duckdb

src/include/storage/mooncake_schema.hpp

@@ -9,7 +9,7 @@ class Moonlink;
 
 class MooncakeSchema : public SchemaCatalogEntry {
 public:
-	MooncakeSchema(Catalog &catalog, CreateSchemaInfo &info);
+	MooncakeSchema(Catalog &catalog, CreateSchemaInfo &info, uint64_t lsn);
 
 	~MooncakeSchema();
 
@@ -48,6 +48,7 @@ class MooncakeSchema : public SchemaCatalogEntry {
 	void Alter(CatalogTransaction transaction, AlterInfo &info) override;
 
 private:
+	uint64_t lsn;
 	Moonlink &moonlink;
 	mutex lock;
 	case_insensitive_map_t<unique_ptr<MooncakeTable>> tables;

src/include/storage/mooncake_table.hpp

@@ -9,7 +9,8 @@ class Moonlink;
 
 class MooncakeTable : public TableCatalogEntry {
 public:
-	MooncakeTable(Catalog &catalog, SchemaCatalogEntry &schema, CreateTableInfo &info, Moonlink &moonlink);
+	MooncakeTable(Catalog &catalog, SchemaCatalogEntry &schema, CreateTableInfo &info, uint64_t lsn,
+	              Moonlink &moonlink);
 
 	~MooncakeTable();
 
@@ -22,6 +23,7 @@ class MooncakeTable : public TableCatalogEntry {
 	MooncakeTableMetadata &GetTableMetadata();
 
 private:
+	uint64_t lsn;
 	Moonlink &moonlink;
 	mutex lock;
 	unique_ptr<MooncakeTableMetadata> metadata;

src/include/storage/mooncake_table_metadata.hpp

@@ -8,7 +8,7 @@ class DeleteFilter;
 
 class MooncakeTableMetadata {
 public:
-	MooncakeTableMetadata(Moonlink &moonlink, const string &schema, const string &table);
+	MooncakeTableMetadata(Moonlink &moonlink, const string &schema, const string &table, uint64_t lsn);
 
 	~MooncakeTableMetadata();
 

src/include/storage/mooncake_transaction.hpp

@@ -14,6 +14,7 @@ class MooncakeTransaction : public Transaction {
 
 private:
 	Catalog &catalog;
+	uint64_t lsn;
 	mutex lock;
 	case_insensitive_map_t<unique_ptr<SchemaCatalogEntry>> schemas;
 };

src/mooncake_extension.cpp

@@ -1,14 +1,25 @@
 #define DUCKDB_EXTENSION_MAIN
 
+#include "duckdb/main/connection.hpp"
 #include "duckdb/main/database.hpp"
 #include "mooncake_extension.hpp"
+#include "pgmooncake.hpp"
 #include "storage/mooncake_storage.hpp"
 
 namespace duckdb {
 
 void MooncakeExtension::Load(DuckDB &db) {
 	auto &config = DBConfig::GetConfig(*db.instance);
 	config.storage_extensions["mooncake"] = make_uniq<MooncakeStorageExtension>();
+
+	string init_query = Pgmooncake::GetInitQuery();
+	if (!init_query.empty()) {
+		Connection connection(db);
+		auto res = connection.Query(init_query);
+		if (res->HasError()) {
+			res->ThrowError();
+		}
+	}
 }
 
 string MooncakeExtension::Name() {

src/moonlink/moonlink.cpp

@@ -11,10 +11,10 @@ DataPtr Moonlink::GetTableSchema(const string &schema, const string &table) {
 	return DataPtr(moonlink_get_table_schema(stream.get(), database.c_str(), schema.c_str(), table.c_str()).Unwrap());
 }
 
-DataPtr Moonlink::ScanTableBegin(const string &schema, const string &table) {
+DataPtr Moonlink::ScanTableBegin(const string &schema, const string &table, uint64_t lsn) {
 	lock_guard<mutex> guard(lock);
 	return DataPtr(
-	    moonlink_scan_table_begin(stream.get(), database.c_str(), schema.c_str(), table.c_str(), 0).Unwrap());
+	    moonlink_scan_table_begin(stream.get(), database.c_str(), schema.c_str(), table.c_str(), lsn).Unwrap());
 }
 
 void Moonlink::ScanTableEnd(const string &schema, const string &table) {

src/storage/mooncake_schema.cpp

@@ -9,8 +9,8 @@
 
 namespace duckdb {
 
-MooncakeSchema::MooncakeSchema(Catalog &catalog, CreateSchemaInfo &info)
-    : SchemaCatalogEntry(catalog, info), moonlink(catalog.Cast<MooncakeCatalog>().GetMoonlink()) {
+MooncakeSchema::MooncakeSchema(Catalog &catalog, CreateSchemaInfo &info, uint64_t lsn)
+    : SchemaCatalogEntry(catalog, info), lsn(lsn), moonlink(catalog.Cast<MooncakeCatalog>().GetMoonlink()) {
 }
 
 MooncakeSchema::~MooncakeSchema() = default;
@@ -96,7 +96,7 @@ optional_ptr<CatalogEntry> MooncakeSchema::LookupEntry(CatalogTransaction transa
 	for (idx_t i = 0; i < names.size(); i++) {
 		table_info.columns.AddColumn(ColumnDefinition(names[i], return_types[i]));
 	}
-	tables[table_name] = make_uniq<MooncakeTable>(catalog, *this, table_info, moonlink);
+	tables[table_name] = make_uniq<MooncakeTable>(catalog, *this, table_info, lsn, moonlink);
 	return *tables[table_name];
 }