vortex_version() macro

Signed-off-by: Mikhail Kot <mikhail@spiraldb.com>

Author: myrrc

vortex-duckdb/build.rs

@@ -627,6 +627,35 @@ fn cbindgen_rust2c(crate_dir: &Path) {
     }
 }
 
+fn git(crate_dir: &Path, args: &[&str]) -> Option<String> {
+    let output = Command::new("git")
+        .arg("-C")
+        .arg(crate_dir)
+        .args(args)
+        .output()
+        .ok()?;
+    if !output.status.success() {
+        return None;
+    }
+    let stdout = String::from_utf8(output.stdout).ok()?;
+    let stdout = stdout.trim();
+    (!stdout.is_empty()).then(|| stdout.to_owned())
+}
+
+fn vortex_version(crate_dir: &Path) {
+    println!("cargo:rerun-if-env-changed=VORTEX_VERSION");
+    let version = env::var("VORTEX_VERSION")
+        .ok()
+        .filter(|version| !version.is_empty())
+        // If this commit belongs to a tag
+        .or_else(|| git(crate_dir, &["describe", "--tags", "--exact-match", "HEAD"]))
+        // If this commit doesn't belong to a tag
+        .or_else(|| git(crate_dir, &["rev-parse", "HEAD"]))
+        // No version, can't build
+        .unwrap();
+    println!("cargo:rustc-env=VORTEX_VERSION={version}");
+}
+
 fn main() {
     println!("cargo:rerun-if-changed=cpp/include");
     println!("cargo:rerun-if-changed=patches");
@@ -653,6 +682,7 @@ fn main() {
     // in vortex's CI.
 
     let crate_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap());
+    vortex_version(&crate_dir);
     if let Some(source_dir) = env::var_os("DUCKDB_SOURCE_DIR") {
         let source_dir = PathBuf::from(source_dir);
         let duckdb_include_dir = source_dir.join("src").join("include");

vortex-duckdb/cpp/include/table_function.h

@@ -76,6 +76,7 @@ typedef struct {
 } duckdb_vx_written_column_statistics;
 
 duckdb_state duckdb_vx_register_table_functions(duckdb_database ffi_db);
+duckdb_state duckdb_vx_register_version_function(duckdb_database ffi_db, const char *version);
 
 typedef struct duckdb_vx_agg_input_ *duckdb_vx_agg_input;
 idx_t duckdb_vx_aggregate_len(duckdb_vx_agg_input ffi);

vortex-duckdb/cpp/table_function.cpp

@@ -12,7 +12,9 @@
 
 #include "duckdb.h"
 #include "duckdb/catalog/catalog.hpp"
+#include "duckdb/catalog/default/default_functions.hpp"
 #include "duckdb/common/insertion_order_preserving_map.hpp"
+#include "duckdb/parser/keyword_helper.hpp"
 #include "duckdb/common/multi_file/multi_file_reader.hpp"
 #include "duckdb/function/table_function.hpp"
 #include "duckdb/main/capi/capi_internal.hpp"
@@ -220,6 +222,32 @@ duckdb_state register_table_function(DatabaseInstance &db, LogicalType parameter
     return DuckDBSuccess;
 }
 
+extern "C" duckdb_state duckdb_vx_register_version_function(duckdb_database ffi_db, const char *version) {
+    D_ASSERT(ffi_db);
+    D_ASSERT(version);
+    const DatabaseWrapper &wrapper = *reinterpret_cast<DatabaseWrapper *>(ffi_db);
+    DatabaseInstance &db = *wrapper.database->instance;
+
+    const string quoted = KeywordHelper::WriteQuoted(version);
+
+    const DefaultMacro macro {DEFAULT_SCHEMA,
+                              "vortex_version",
+                              {nullptr},
+                              {{nullptr, nullptr}},
+                              quoted.c_str()};
+    try {
+        auto info = DefaultFunctionGenerator::CreateInternalMacroInfo(macro);
+        auto &system_catalog = Catalog::GetSystemCatalog(db);
+        auto data = CatalogTransaction::GetSystemTransaction(db);
+        system_catalog.CreateFunction(data, *info);
+    } catch (const std::exception &e) {
+        ErrorData data(e);
+        DUCKDB_LOG_ERROR(db, "Failed to create the vortex_version macro:\t" + data.Message());
+        return DuckDBError;
+    }
+    return DuckDBSuccess;
+}
+
 extern "C" duckdb_state duckdb_vx_register_table_functions(duckdb_database ffi_db) {
     D_ASSERT(ffi_db);
     const DatabaseWrapper &wrapper = *reinterpret_cast<DatabaseWrapper *>(ffi_db);

vortex-duckdb/src/duckdb/database.rs

@@ -62,6 +62,16 @@ impl DatabaseRef {
         Ok(())
     }
 
+    pub fn register_version_function(&self, version: &str) -> VortexResult<()> {
+        let version = CString::new(version)
+            .map_err(|_| vortex_err!("Invalid version: string contains null bytes"))?;
+        duckdb_try!(
+            unsafe { cpp::duckdb_vx_register_version_function(self.as_ptr(), version.as_ptr()) },
+            "Failed to register vortex_version function"
+        );
+        Ok(())
+    }
+
     pub fn register_optimizer_extension(&self) -> VortexResult<()> {
         duckdb_try!(
             unsafe { cpp::duckdb_vx_optimizer_extension_register(self.as_ptr()) },

vortex-duckdb/src/e2e_test/vortex_scan_test.rs

@@ -201,6 +201,19 @@ fn test_scan_function_registration() {
     assert_eq!(string, "vortex_scan");
 }
 
+#[test]
+fn test_vortex_version() -> Result<()> {
+    let conn = database_connection();
+    let query = format!(
+        "SELECT (vortex_version() = '{}')::INT",
+        env!("VORTEX_VERSION")
+    );
+    let result = conn.query(&query)?;
+    let chunk = result.into_iter().next().unwrap();
+    assert_eq!(chunk.get_vector(0).as_slice_with_len::<i32>(1), [1]);
+    Ok(())
+}
+
 #[test]
 fn test_vortex_scan_strings() {
     let file = RUNTIME.block_on(async {

vortex-duckdb/src/lib.rs

@@ -78,6 +78,7 @@ fn init_tracing() {
 /// separately (e.g., before creating connections), call `register_extension_options` first.
 pub fn initialize(db: &DatabaseRef) -> VortexResult<()> {
     db.register_table_functions()?;
+    db.register_version_function(env!("VORTEX_VERSION"))?;
     db.register_optimizer_extension()?;
     db.register_copy_function()
 }

vortex-sqllogictest/slt/duckdb/version.slt

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+include ../setup.slt.no
+
+query B
+SELECT length(vortex_version()) > 0;
+----
+true