WIP cache schema function

Kayrnt · Kayrnt · commit 1790f3a8ad43 · 2023-11-20T09:58:40.000+01:00
diff --git a/duckdb b/duckdb
@@ -1 +1 @@
-Subproject commit 3c695d7ba94d95d9facee48d395f46ed0bd72b46
+Subproject commit 401c8061c6ece35949cac58c7770cc755710ca86
diff --git a/src/include/mysql_scanner.hpp b/src/include/mysql_scanner.hpp
@@ -39,6 +39,11 @@ class MySQLScanFunction : public TableFunction {
 	MySQLScanFunction();
 };
 
+class MySQLCacheSchemaFunction : public TableFunction {
+public:
+	MySQLCacheSchemaFunction();
+};
+
 class MySQLClearCacheFunction : public TableFunction {
 public:
 	MySQLClearCacheFunction();
diff --git a/src/include/storage/mysql_catalog.hpp b/src/include/storage/mysql_catalog.hpp
@@ -55,6 +55,7 @@ class MySQLCatalog : public Catalog {
 	bool InMemory() override;
 	string GetDBPath() override;
 
+	void CacheSchema(ClientContext &context);
 	void ClearCache();
 
 private:
diff --git a/src/include/storage/mysql_catalog_set.hpp b/src/include/storage/mysql_catalog_set.hpp
@@ -24,6 +24,7 @@ class MySQLCatalogSet {
 	virtual void DropEntry(ClientContext &context, DropInfo &info);
 	void Scan(ClientContext &context, const std::function<void(CatalogEntry &)> &callback);
 	optional_ptr<CatalogEntry> CreateEntry(unique_ptr<CatalogEntry> entry);
+	void CacheEntries(ClientContext &context);
 	void ClearEntries();
 
 protected:
diff --git a/src/include/storage/mysql_schema_entries_cache.hpp b/src/include/storage/mysql_schema_entries_cache.hpp
@@ -0,0 +1,37 @@
+#pragma once
+
+#include "duckdb.hpp"
+#ifndef DUCKDB_AMALGAMATION
+#include "duckdb/storage/object_cache.hpp"
+#endif
+//#include "mysql_schema_entry.hpp"
+
+namespace duckdb {
+
+//! MysqlSchemaEntriesCache
+class MysqlSchemaEntriesCache : public ObjectCacheEntry {
+public:
+	MysqlSchemaEntriesCache() : entries(nullptr) {
+	}
+	MysqlSchemaEntriesCache(unique_ptr<vector<string>> mysql_schema_entries, time_t r_time)
+	    : entries(std::move(mysql_schema_entries)), read_time(r_time) {
+	}
+
+	~MysqlSchemaEntriesCache() override = default;
+
+	//! Mysql schema entries
+	unique_ptr<vector<string>> entries;
+
+	//! read time
+	time_t read_time;
+
+public:
+	static string ObjectType() {
+		return "mysql_schema_entries";
+	}
+
+	string GetObjectType() override {
+		return ObjectType();
+	}
+};
+} // namespace duckdb
diff --git a/src/include/storage/mysql_schema_set.hpp b/src/include/storage/mysql_schema_set.hpp
@@ -23,6 +23,10 @@ class MySQLSchemaSet : public MySQLCatalogSet {
 
 protected:
 	void LoadEntries(ClientContext &context) override;
+  void PutInCache(ClientContext &context, vector<string> &schemas);
+  vector<string> LoadSchemasFromMysqlInformationSchema(ClientContext &context);
 };
 
+extern const string MYSQL_SCHEMA_ENTRIES_CACHE_KEY;
+
 } // namespace duckdb
diff --git a/src/storage/CMakeLists.txt b/src/storage/CMakeLists.txt
@@ -1,5 +1,6 @@
 add_library(
   mysql_ext_storage OBJECT
+  mysql_cache_schema.cpp
   mysql_catalog.cpp
   mysql_catalog_set.cpp
   mysql_clear_cache.cpp
diff --git a/src/storage/mysql_cache_schema.cpp b/src/storage/mysql_cache_schema.cpp
@@ -0,0 +1,44 @@
+#include "duckdb.hpp"
+
+#include "duckdb/parser/parsed_data/create_table_function_info.hpp"
+#include "mysql_scanner.hpp"
+#include "duckdb/main/database_manager.hpp"
+#include "duckdb/main/attached_database.hpp"
+#include "storage/mysql_catalog.hpp"
+
+namespace duckdb {
+
+struct CacheSchemaFunctionData : public TableFunctionData {
+	bool finished = false;
+};
+
+static unique_ptr<FunctionData> CacheSchemaBind(ClientContext &context, TableFunctionBindInput &input,
+                                               vector<LogicalType> &return_types, vector<string> &names) {
+
+	auto result = make_uniq<CacheSchemaFunctionData>();
+	return_types.push_back(LogicalType::BOOLEAN);
+	names.emplace_back("Success");
+	return std::move(result);
+}
+
+static void CacheSchemaFunction(ClientContext &context, TableFunctionInput &data_p, DataChunk &output) {
+	auto &data = data_p.bind_data->CastNoConst<CacheSchemaFunctionData>();
+	if (data.finished) {
+		return;
+	}
+	auto databases = DatabaseManager::Get(context).GetDatabases(context);
+	for (auto &db_ref : databases) {
+		auto &db = db_ref.get();
+		auto &catalog = db.GetCatalog();
+		if (catalog.GetCatalogType() != "mysql") {
+			continue;
+		}
+		catalog.Cast<MySQLCatalog>().CacheSchema(context);
+	}
+	data.finished = true;
+}
+
+MySQLCacheSchemaFunction::MySQLCacheSchemaFunction()
+    : TableFunction("mysql_cache_schema", {}, CacheSchemaFunction, CacheSchemaBind) {
+}
+} // namespace duckdb
diff --git a/src/storage/mysql_catalog.cpp b/src/storage/mysql_catalog.cpp
@@ -92,6 +92,10 @@ WHERE table_schema = ${SCHEMA_NAME};
 	return size;
 }
 
+void MySQLCatalog::CacheSchema(ClientContext &context) {
+	schemas.CacheEntries(context);
+}
+
 void MySQLCatalog::ClearCache() {
 	schemas.ClearEntries();
 }
diff --git a/src/storage/mysql_catalog_set.cpp b/src/storage/mysql_catalog_set.cpp
@@ -64,6 +64,13 @@ optional_ptr<CatalogEntry> MySQLCatalogSet::CreateEntry(unique_ptr<CatalogEntry>
 	return result;
 }
 
+void MySQLCatalogSet::CacheEntries(ClientContext &context) {
+	if (!is_loaded) {
+		is_loaded = true;
+		LoadEntries(context);
+	}
+}
+
 void MySQLCatalogSet::ClearEntries() {
 	entries.clear();
 	is_loaded = false;
diff --git a/src/storage/mysql_schema_set.cpp b/src/storage/mysql_schema_set.cpp
@@ -1,25 +1,61 @@
 #include "storage/mysql_schema_set.hpp"
 #include "storage/mysql_transaction.hpp"
 #include "duckdb/parser/parsed_data/create_schema_info.hpp"
+#include "duckdb/storage/object_cache.hpp"
+#include "storage/mysql_schema_entries_cache.hpp"
+#include <iostream>
 
 namespace duckdb {
 
+const std::string MYSQL_SCHEMA_ENTRIES_CACHE_KEY = "mysql_schema_entries_cache";
+
 MySQLSchemaSet::MySQLSchemaSet(Catalog &catalog) : MySQLCatalogSet(catalog) {
 }
 
 void MySQLSchemaSet::LoadEntries(ClientContext &context) {
+	auto db = context.db;
+	auto& object_cache = db->GetObjectCache();
+	auto schema_cache = object_cache.Get<MysqlSchemaEntriesCache>(MYSQL_SCHEMA_ENTRIES_CACHE_KEY);
+	if(schema_cache != nullptr) {
+		//std::cout << "load from cache" << std::endl;
+		auto mysql_schema_entries = std::move(schema_cache->entries);
+		for(auto& schema_name : *mysql_schema_entries) {
+			auto schema = make_uniq<MySQLSchemaEntry>(catalog, schema_name);
+			CreateEntry(std::move(schema));
+		}
+	} else {
+		//std::cout << "load from mysql" << std::endl;
+		auto schemas = LoadSchemasFromMysqlInformationSchema(context);
+		for(auto& schema : schemas) {
+			auto schema_entry = make_uniq<MySQLSchemaEntry>(catalog, schema);
+			CreateEntry(std::move(schema_entry));
+		}
+		PutInCache(context, schemas);
+	}
+}
+
+void MySQLSchemaSet::PutInCache(ClientContext &context, vector<string> &schemas) {
+	auto db = context.db;
+	auto& object_cache = db->GetObjectCache();
+	auto mysql_schema_entries = make_shared<MysqlSchemaEntriesCache>(make_uniq<vector<string>>(schemas), time(nullptr));
+	object_cache.Put(MYSQL_SCHEMA_ENTRIES_CACHE_KEY, mysql_schema_entries);
+}
+
+vector<string> MySQLSchemaSet::LoadSchemasFromMysqlInformationSchema(ClientContext &context) {
 	auto query = R"(
 SELECT schema_name
 FROM information_schema.schemata;
 )";
-
+	
 	auto &transaction = MySQLTransaction::Get(context, catalog);
 	auto result = transaction.Query(query);
+	
+	auto schemas = vector<string>();
 	while (result->Next()) {
 		auto schema_name = result->GetString(0);
-		auto schema = make_uniq<MySQLSchemaEntry>(catalog, schema_name);
-		CreateEntry(std::move(schema));
+		schemas.push_back(schema_name);
 	}
+	return schemas;
 }
 
 optional_ptr<CatalogEntry> MySQLSchemaSet::CreateSchema(ClientContext &context, CreateSchemaInfo &info) {

Original file line number	Diff line number	Diff line change
`@@ -92,6 +92,10 @@ WHERE table_schema = ${SCHEMA_NAME};`
`92`	`92`	`return size;`
`93`	`93`	`}`
`94`	`94`
	`95`	`+void MySQLCatalog::CacheSchema(ClientContext &context) {`
	`96`	`+ schemas.CacheEntries(context);`
	`97`	`+}`
	`98`	`+`
`95`	`99`	`void MySQLCatalog::ClearCache() {`
`96`	`100`	`schemas.ClearEntries();`
`97`	`101`	`}`