feat(uast): add Go/Rust parsing and projections

Author: marcialwushu

Cargo.lock

@@ -5,6 +5,7 @@ members = [
     "crates/gitbase-git",
     "crates/gitbase-loader",
     "crates/gitbase-pgwire",
+    "crates/gitbase-uast",
 ]
 resolver = "2"
 
@@ -19,10 +20,15 @@ async-trait = "0.1"
 clap = { version = "4", features = ["derive", "env"] }
 gix = "0.66"
 pgwire = "0.38"
+serde = { version = "1", features = ["derive"] }
+serde_json = "1"
 sha1 = "0.10"
 sqlx = { version = "0.8", features = ["postgres", "runtime-tokio", "migrate", "macros"] }
 tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal", "net"] }
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] }
 futures = "0.3"
+tree-sitter = "0.24"
+tree-sitter-go = "0.23"
+tree-sitter-rust = "0.23"
 walkdir = "2"

Cargo.toml

@@ -81,6 +81,21 @@ enum Commands {
         #[arg(long, env = "GITBASE_DB_MAX_CONNECTIONS", default_value_t = 5)]
         max_connections: u32,
     },
+
+    /// Build UAST cache and projections
+    Uast {
+        /// PostgreSQL connection string
+        #[arg(long, env = "DATABASE_URL")]
+        database_url: String,
+
+        /// Maximum database connections
+        #[arg(long, env = "GITBASE_DB_MAX_CONNECTIONS", default_value_t = 5)]
+        max_connections: u32,
+
+        /// Limit number of blobs to process
+        #[arg(long, env = "GITBASE_UAST_LIMIT")]
+        limit: Option<i64>,
+    },
 }
 
 #[tokio::main]
@@ -154,6 +169,29 @@ async fn main() -> Result<()> {
             gitbase_db::health_check(&pool).await?;
             tracing::info!("health check passed");
         }
+        Commands::Uast {
+            database_url,
+            max_connections,
+            limit,
+        } => {
+            let pool = gitbase_db::connect(&database_url, max_connections).await?;
+            gitbase_db::health_check(&pool).await?;
+            tracing::info!("health check passed");
+
+            let report = gitbase_loader::index_uast(
+                &pool,
+                &gitbase_loader::UastIndexConfig { max_candidates: limit },
+            )
+            .await?;
+
+            tracing::info!(
+                parsed = report.parsed,
+                skipped_cached = report.skipped_cached,
+                skipped_missing_content = report.skipped_missing_content,
+                skipped_unsupported_language = report.skipped_unsupported_language,
+                "uast indexing completed"
+            );
+        }
     }
 
     Ok(())

crates/gitbase-cli/src/main.rs

@@ -5,6 +5,7 @@ edition.workspace = true
 
 [dependencies]
 anyhow.workspace = true
+serde_json.workspace = true
 sqlx.workspace = true
 tokio.workspace = true
 tracing.workspace = true

crates/gitbase-db/Cargo.toml

@@ -99,3 +99,12 @@ pub async fn fetch_blob_locations(
         })
         .collect())
 }
+
+pub async fn fetch_blob_content(pool: &PgPool, hash: &str) -> Result<Option<Vec<u8>>> {
+    let row = sqlx::query("SELECT content FROM gitbase.blobs WHERE hash = $1")
+        .bind(hash)
+        .fetch_optional(pool)
+        .await?;
+
+    Ok(row.and_then(|row| row.get::<Option<Vec<u8>>, _>("content")))
+}

crates/gitbase-db/src/blobs.rs

@@ -4,6 +4,7 @@ use sqlx::PgPool;
 
 pub mod blobs;
 pub mod metadata;
+pub mod uast;
 
 /// Create a connection pool to PostgreSQL and run migrations.
 pub async fn connect(database_url: &str, max_connections: u32) -> Result<PgPool> {

crates/gitbase-db/src/lib.rs

@@ -0,0 +1,126 @@
+use anyhow::Result;
+use serde_json::Value;
+use sqlx::{PgPool, Row};
+
+#[derive(Debug, Clone)]
+pub struct UastCandidate {
+    pub blob_hash: String,
+    pub path: String,
+}
+
+#[derive(Debug, Clone)]
+pub struct UastFunctionRecord {
+    pub blob_hash: String,
+    pub name: String,
+    pub start_line: Option<i32>,
+    pub end_line: Option<i32>,
+    pub signature: Option<String>,
+}
+
+#[derive(Debug, Clone)]
+pub struct UastImportRecord {
+    pub blob_hash: String,
+    pub source: String,
+    pub target: Option<String>,
+}
+
+pub async fn uast_cache_exists(pool: &PgPool, blob_hash: &str) -> Result<bool> {
+    let exists = sqlx::query_scalar::<_, bool>(
+        "SELECT EXISTS(SELECT 1 FROM gitbase.uast_cache WHERE blob_hash = $1)",
+    )
+    .bind(blob_hash)
+    .fetch_one(pool)
+    .await?;
+    Ok(exists)
+}
+
+pub async fn fetch_uast_candidates(pool: &PgPool, limit: Option<i64>) -> Result<Vec<UastCandidate>> {
+    let base_query =
+        "SELECT DISTINCT f.blob_hash, f.path FROM gitbase.files f\
+         JOIN gitbase.blobs b ON b.hash = f.blob_hash\
+         WHERE b.content IS NOT NULL";
+
+    let rows = if let Some(limit) = limit {
+        sqlx::query(&format!("{base_query} LIMIT $1"))
+            .bind(limit)
+            .fetch_all(pool)
+            .await?
+    } else {
+        sqlx::query(base_query).fetch_all(pool).await?
+    };
+
+    Ok(rows
+        .into_iter()
+        .map(|row| UastCandidate {
+            blob_hash: row.get("blob_hash"),
+            path: row.get("path"),
+        })
+        .collect())
+}
+
+pub async fn upsert_uast_cache(
+    pool: &PgPool,
+    blob_hash: &str,
+    language: &str,
+    uast: &Value,
+) -> Result<()> {
+    sqlx::query(
+        "INSERT INTO gitbase.uast_cache (blob_hash, language, uast, generated_at)\
+         VALUES ($1, $2, $3, now())\
+         ON CONFLICT (blob_hash) DO UPDATE\
+         SET language = EXCLUDED.language,\
+             uast = EXCLUDED.uast,\
+             generated_at = now()",
+    )
+    .bind(blob_hash)
+    .bind(language)
+    .bind(uast)
+    .execute(pool)
+    .await?;
+
+    Ok(())
+}
+
+pub async fn clear_uast_projections(pool: &PgPool, blob_hash: &str) -> Result<()> {
+    sqlx::query("DELETE FROM gitbase.uast_functions WHERE blob_hash = $1")
+        .bind(blob_hash)
+        .execute(pool)
+        .await?;
+    sqlx::query("DELETE FROM gitbase.uast_imports WHERE blob_hash = $1")
+        .bind(blob_hash)
+        .execute(pool)
+        .await?;
+    Ok(())
+}
+
+pub async fn insert_uast_function(pool: &PgPool, record: &UastFunctionRecord) -> Result<()> {
+    sqlx::query(
+        "INSERT INTO gitbase.uast_functions (blob_hash, name, start_line, end_line, signature)\
+         VALUES ($1, $2, $3, $4, $5)\
+         ON CONFLICT (blob_hash, name, start_line) DO UPDATE\
+         SET end_line = EXCLUDED.end_line,\
+             signature = EXCLUDED.signature",
+    )
+    .bind(&record.blob_hash)
+    .bind(&record.name)
+    .bind(record.start_line)
+    .bind(record.end_line)
+    .bind(&record.signature)
+    .execute(pool)
+    .await?;
+    Ok(())
+}
+
+pub async fn insert_uast_import(pool: &PgPool, record: &UastImportRecord) -> Result<()> {
+    sqlx::query(
+        "INSERT INTO gitbase.uast_imports (blob_hash, source, target)\
+         VALUES ($1, $2, $3)\
+         ON CONFLICT (blob_hash, source, target) DO NOTHING",
+    )
+    .bind(&record.blob_hash)
+    .bind(&record.source)
+    .bind(&record.target)
+    .execute(pool)
+    .await?;
+    Ok(())
+}

crates/gitbase-db/src/uast.rs

@@ -7,6 +7,8 @@ edition.workspace = true
 anyhow.workspace = true
 gitbase-db = { path = "../gitbase-db" }
 gitbase-git = { path = "../gitbase-git" }
+gitbase-uast = { path = "../gitbase-uast" }
+serde_json.workspace = true
 sqlx.workspace = true
 tracing.workspace = true