store: Add resumable entity data import for restore

Add import_data() to the restore module that reads Parquet chunks
and inserts entity data into PostgreSQL tables. Supports resumability
by checking max(vid) in each table and skipping already-imported rows.

Entity tables use InsertQuery::for_restore() for efficient batch
inserts. The data_sources$ table uses raw SQL with bind parameters
since it has a fixed schema outside the Layout.

Author: lutter

store/postgres/src/relational/restore.rs

@@ -10,21 +10,29 @@ use std::fs;
 use std::path::Path;
 use std::sync::Arc;
 
-use diesel_async::SimpleAsyncConnection;
-
+use diesel::sql_types::{BigInt, Binary, Integer, Nullable, Text};
+use diesel_async::{RunQueryDsl, SimpleAsyncConnection};
 use graph::blockchain::BlockHash;
 use graph::data::subgraph::schema::{DeploymentCreate, SubgraphManifestEntity};
-use graph::prelude::{BlockPtr as GraphBlockPtr, StoreError};
+use graph::prelude::{info, BlockPtr as GraphBlockPtr, Logger, StoreError};
 use graph::schema::{EntityType, InputSchema};
 use graph::semver::Version;
 
 use crate::deployment::create_deployment;
 use crate::dynds::DataSourcesTable;
+use crate::parquet::convert::{
+    record_batch_to_data_source_rows, record_batch_to_restore_rows, DataSourceRestoreRow,
+};
+use crate::parquet::reader::read_batches;
 use crate::primary::Site;
-use crate::relational::dump::Metadata;
-use crate::relational::Layout;
+use crate::relational::dump::{Metadata, TableInfo};
+use crate::relational::{Layout, Table};
+use crate::relational_queries::InsertQuery;
+use crate::vid_batcher::VidRange;
 use crate::AsyncPgConnection;
 
+const DATA_SOURCES_TABLE: &str = "data_sources$";
+
 /// Convert a dump `BlockPtr` (hex hash string) to a graph `BlockPtr`.
 fn to_graph_block_ptr(bp: &super::dump::BlockPtr) -> Result<GraphBlockPtr, StoreError> {
     let hash = BlockHash::try_from(bp.hash.as_str())
@@ -35,6 +43,167 @@ fn to_graph_block_ptr(bp: &super::dump::BlockPtr) -> Result<GraphBlockPtr, Store
     })
 }
 
+/// Query the current max(vid) for a table. Returns -1 if the table is empty.
+async fn current_max_vid(
+    conn: &mut AsyncPgConnection,
+    qualified_name: &str,
+) -> Result<i64, StoreError> {
+    let query = format!(
+        "select coalesce(min(vid), 0)::int8 as min_vid, \
+                coalesce(max(vid), -1)::int8 as max_vid \
+           from {}",
+        qualified_name
+    );
+    let range: VidRange = diesel::sql_query(&query)
+        .get_result(conn)
+        .await
+        .map_err(StoreError::from)?;
+    Ok(range.max)
+}
+
+/// Import a single entity table from Parquet chunks.
+///
+/// Supports resumability: checks the current max(vid) in the DB table
+/// and skips already-imported rows.
+async fn import_entity_table(
+    conn: &mut AsyncPgConnection,
+    table: &Table,
+    table_info: &TableInfo,
+    dir: &Path,
+    logger: &Logger,
+) -> Result<usize, StoreError> {
+    if table_info.chunks.is_empty() || table_info.max_vid < 0 {
+        return Ok(0);
+    }
+
+    let max_vid_db = current_max_vid(conn, table.qualified_name.as_str()).await?;
+    if max_vid_db >= table_info.max_vid {
+        info!(
+            logger,
+            "Table {} already fully restored, skipping",
+            table.object.as_str()
+        );
+        return Ok(0);
+    }
+
+    let chunk_size = InsertQuery::chunk_size(table);
+    let mut total_inserted = 0usize;
+
+    for chunk_info in &table_info.chunks {
+        // Skip chunks that are fully imported
+        if chunk_info.max_vid <= max_vid_db {
+            continue;
+        }
+
+        let chunk_path = dir.join(&chunk_info.file);
+        let batches = read_batches(&chunk_path)?;
+
+        for batch in &batches {
+            let mut rows = record_batch_to_restore_rows(batch, table)?;
+
+            // Filter out already-imported rows (for boundary chunks on resume)
+            if max_vid_db >= 0 {
+                rows.retain(|row| row.vid > max_vid_db);
+            }
+
+            if rows.is_empty() {
+                continue;
+            }
+
+            // Split into InsertQuery-sized chunks and execute
+            for chunk in rows.chunks(chunk_size) {
+                InsertQuery::for_restore(table, chunk)?
+                    .execute(conn)
+                    .await?;
+                total_inserted += chunk.len();
+            }
+        }
+    }
+
+    info!(
+        logger,
+        "Restored {} rows into {}",
+        total_inserted,
+        table.object.as_str()
+    );
+    Ok(total_inserted)
+}
+
+/// Insert a single data_sources$ row via raw SQL.
+async fn insert_data_source_row(
+    conn: &mut AsyncPgConnection,
+    qualified_table: &str,
+    row: &DataSourceRestoreRow,
+) -> Result<(), StoreError> {
+    let query = format!(
+        "INSERT INTO {} (vid, block_range, causality_region, manifest_idx, \
+                parent, id, param, context, done_at) \
+         VALUES ($1, int4range($2, $3), $4, $5, $6, $7, $8, $9::jsonb, $10)",
+        qualified_table,
+    );
+    diesel::sql_query(&query)
+        .bind::<BigInt, _>(row.vid)
+        .bind::<Integer, _>(row.block_range_start)
+        .bind::<Nullable<Integer>, _>(row.block_range_end)
+        .bind::<Integer, _>(row.causality_region)
+        .bind::<Integer, _>(row.manifest_idx)
+        .bind::<Nullable<Integer>, _>(row.parent)
+        .bind::<Nullable<Binary>, _>(row.id.as_deref())
+        .bind::<Nullable<Binary>, _>(row.param.as_deref())
+        .bind::<Nullable<Text>, _>(row.context.as_deref())
+        .bind::<Nullable<Integer>, _>(row.done_at)
+        .execute(conn)
+        .await
+        .map_err(StoreError::from)?;
+    Ok(())
+}
+
+/// Import the `data_sources$` table from Parquet chunks.
+async fn import_data_sources(
+    conn: &mut AsyncPgConnection,
+    namespace: &str,
+    table_info: &TableInfo,
+    dir: &Path,
+    logger: &Logger,
+) -> Result<usize, StoreError> {
+    if table_info.chunks.is_empty() || table_info.max_vid < 0 {
+        return Ok(0);
+    }
+
+    let qualified = format!("\"{}\".\"{DATA_SOURCES_TABLE}\"", namespace);
+    let max_vid_db = current_max_vid(conn, &qualified).await?;
+    if max_vid_db >= table_info.max_vid {
+        info!(logger, "data_sources$ already fully restored, skipping");
+        return Ok(0);
+    }
+
+    let mut total_inserted = 0usize;
+
+    for chunk_info in &table_info.chunks {
+        if chunk_info.max_vid <= max_vid_db {
+            continue;
+        }
+
+        let chunk_path = dir.join(&chunk_info.file);
+        let batches = read_batches(&chunk_path)?;
+
+        for batch in &batches {
+            let rows = record_batch_to_data_source_rows(batch)?;
+
+            for row in &rows {
+                if max_vid_db >= 0 && row.vid <= max_vid_db {
+                    continue;
+                }
+                insert_data_source_row(conn, &qualified, row).await?;
+                total_inserted += 1;
+            }
+        }
+    }
+
+    info!(logger, "Restored {} data_sources$ rows", total_inserted);
+    Ok(total_inserted)
+}
+
 /// Restore a subgraph deployment's schema and metadata from a dump
 /// directory.
 ///
@@ -132,11 +301,57 @@ pub async fn create_schema(
     .await?;
 
     // 8. Create data_sources$ table if present in dump
-    if metadata.tables.contains_key("data_sources$") {
+    if metadata.tables.contains_key(DATA_SOURCES_TABLE) {
         let ds_table = DataSourcesTable::new(site.namespace.clone());
         let ddl = ds_table.as_ddl();
         conn.batch_execute(&ddl).await?;
     }
 
     Ok(layout)
 }
+
+/// Import entity data and data_sources$ from Parquet files into the
+/// database tables created by `create_schema`.
+///
+/// This is resumable: if interrupted, it can be called again and will
+/// skip already-imported rows by checking the current max(vid) in each
+/// table.
+#[allow(dead_code)]
+pub async fn import_data(
+    conn: &mut AsyncPgConnection,
+    layout: &Layout,
+    metadata: &Metadata,
+    dir: &Path,
+    logger: &Logger,
+) -> Result<(), StoreError> {
+    // Import entity tables (sorted by name for determinism)
+    let mut table_names: Vec<_> = metadata
+        .tables
+        .keys()
+        .filter(|name| name.as_str() != DATA_SOURCES_TABLE)
+        .collect();
+    table_names.sort();
+
+    for table_name in table_names {
+        let table_info = &metadata.tables[table_name];
+        let table = layout
+            .tables
+            .values()
+            .find(|t| t.object.as_str() == table_name)
+            .ok_or_else(|| {
+                StoreError::InternalError(format!(
+                    "table '{}' from dump not found in layout",
+                    table_name,
+                ))
+            })?;
+        import_entity_table(conn, table, table_info, dir, logger).await?;
+    }
+
+    // Import data_sources$ if present
+    if let Some(ds_info) = metadata.tables.get(DATA_SOURCES_TABLE) {
+        let namespace = layout.site.namespace.as_str();
+        import_data_sources(conn, namespace, ds_info, dir, logger).await?;
+    }
+
+    Ok(())
+}