[SpiceDQ] Improve error messages; Avoid race condition on allocate_initial_partitions. (spiceai#9579)

* make errors better

* Avoid issues when executor requests allocate_initial_partitions earlier than scheduler can handle

* fix lint

* linting

Author: Jeadie

crates/runtime/src/cluster/partition/manager.rs

@@ -40,6 +40,9 @@ pub enum Error {
     #[snafu(display("Partition {partition} not found in table {table}"))]
     PartitionNotFound { table: String, partition: String },
 
+    #[snafu(display("No partition metadata found for table {table}"))]
+    TableMetadataNotFound { table: String },
+
     #[snafu(display("Concurrent modification detected for table {table}"))]
     ConcurrentModification { table: String },
 }
@@ -152,13 +155,10 @@ impl PartitionManager {
 
         loop {
             let now_ms = now_ms()?;
-            let mut metadata =
-                self.get_table_metadata(table)
-                    .await?
-                    .ok_or_else(|| Error::PartitionNotFound {
-                        table: key.clone(),
-                        partition: "any".to_string(),
-                    })?;
+            let mut metadata = self
+                .get_table_metadata(table)
+                .await?
+                .ok_or_else(|| Error::TableMetadataNotFound { table: key.clone() })?;
 
             let mut allocated: Vec<_> = metadata
                 .partitions
@@ -219,13 +219,10 @@ impl PartitionManager {
 
         loop {
             let now_ms = now_ms()?;
-            let mut metadata =
-                self.get_table_metadata(table)
-                    .await?
-                    .ok_or_else(|| Error::PartitionNotFound {
-                        table: key.clone(),
-                        partition: "any".to_string(),
-                    })?;
+            let mut metadata = self
+                .get_table_metadata(table)
+                .await?
+                .ok_or_else(|| Error::TableMetadataNotFound { table: key.clone() })?;
 
             let mut updated = false;
             for partition in &mut metadata.partitions {

crates/runtime/src/cluster/service.rs

@@ -603,8 +603,16 @@ impl ClusterService for ClusterServiceImpl {
         let app_guard = self.app.read().await;
         if let Some(app) = app_guard.as_ref() {
             // Find accelerated datasets with partitioning
-
             for table_ref in super::partition::accelerated_tables(app).keys() {
+                if partition_manager
+                    .get_cached_table_metadata(table_ref)
+                    .is_none()
+                {
+                    tracing::info!(
+                        "No cached partition metadata for table {table_ref}. Scheduler likely has not finished discovering partitions for the table. Will not assign in initial allocation, but will get assigned on future assignments"
+                    );
+                    continue;
+                }
                 match partition_manager
                     .allocate_partitions(table_ref, executor_id, 10)
                     .await

crates/runtime/src/datafusion/refresh_sql.rs

@@ -47,8 +47,13 @@ pub enum Error {
     ))]
     ExpectedSingleSqlStatement { num_statements: usize },
 
-    #[snafu(display("Expected a SQL query starting with SELECT <columns> FROM {expected_table}"))]
-    InvalidSqlStatement { expected_table: TableReference },
+    #[snafu(display(
+        "Expected a SQL query starting with SELECT <columns> FROM {expected_table}. {issue}"
+    ))]
+    InvalidSqlStatement {
+        expected_table: TableReference,
+        issue: String,
+    },
 
     #[snafu(display(
         "Unexpected '{expr}' in the Refresh SQL statement. Rewrite the SQL to only perform WHERE filters, i.e. SELECT col1, col2, col3 FROM {expected_table} WHERE col1 = 'foo'"
@@ -127,7 +132,13 @@ pub fn parse_refresh_sql(
                         )?;
                         ensure!(
                             select.from.len() == 1,
-                            InvalidSqlStatementSnafu { expected_table }
+                            InvalidSqlStatementSnafu {
+                                expected_table,
+                                issue: format!(
+                                    "A single FROM clause with table reference is expected, found {}",
+                                    select.from.len()
+                                )
+                            }
                         );
 
                         ensure_no_expr!(select.cluster_by.is_empty(), "CLUSTER BY", expected_table);
@@ -178,11 +189,23 @@ pub fn parse_refresh_sql(
                                 if TableReference::parse_str(&table_name_with_schema)
                                     != expected_table
                                 {
-                                    return InvalidSqlStatementSnafu { expected_table }.fail();
+                                    return InvalidSqlStatementSnafu {
+                                        expected_table: expected_table.clone(),
+                                        issue: format!(
+                                            "Table name in refresh_sql should be {expected_table}, is {name}",
+                                        ),
+                                    }
+                                    .fail();
                                 }
                             }
                             _ => {
-                                return InvalidSqlStatementSnafu { expected_table }.fail();
+                                return InvalidSqlStatementSnafu {
+                                    expected_table,
+                                    issue:
+                                        "No FROM clause with table reference found in SQL statement"
+                                            .to_string(),
+                                }
+                                .fail();
                             }
                         }
 
@@ -198,12 +221,27 @@ pub fn parse_refresh_sql(
 
                         Ok((refresh_sql, refresh_schema))
                     }
-                    _ => InvalidSqlStatementSnafu { expected_table }.fail()?,
+                    _ => InvalidSqlStatementSnafu {
+                        expected_table,
+                        issue: format!(
+                            "Expected a basic Select SQL query statement, found {}",
+                            query.body
+                        ),
+                    }
+                    .fail()?,
                 }
             }
-            _ => InvalidSqlStatementSnafu { expected_table }.fail()?,
+            _ => InvalidSqlStatementSnafu {
+                expected_table,
+                issue: format!("Expected a Select SQL query statement, found {statement}"),
+            }
+            .fail()?,
         },
-        _ => InvalidSqlStatementSnafu { expected_table }.fail()?,
+        _ => InvalidSqlStatementSnafu {
+            expected_table,
+            issue: format!("Expected a SQL Statement, found {statement}"),
+        }
+        .fail()?,
     }
 }