From 75b0fb58465fc2a1af075a2f6e5086ee0e5b5f8b Mon Sep 17 00:00:00 2001
From: sgrebnov <sergei.grebnov@gmail.com>
Date: Tue, 2 Apr 2024 17:37:48 -0700
Subject: [PATCH] Check connectivity for the Spice.ai endpoint

---
 crates/runtime/src/dataconnector/spiceai.rs | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/crates/runtime/src/dataconnector/spiceai.rs b/crates/runtime/src/dataconnector/spiceai.rs
index 36653cfcbb..01cc9466fa 100644
--- a/crates/runtime/src/dataconnector/spiceai.rs
+++ b/crates/runtime/src/dataconnector/spiceai.rs
@@ -21,6 +21,7 @@ use flight_client::FlightClient;
 use flight_datafusion::FlightTable;
 use futures::StreamExt;
 use futures_core::stream::BoxStream;
+use ns_lookup::verify_endpoint_connection;
 use secrets::Secret;
 use snafu::prelude::*;
 use spicepod::component::dataset::Dataset;
@@ -77,6 +78,11 @@ impl DataConnectorFactory for SpiceAI {
                 .and_then(|params| params.get("endpoint").cloned())
                 .unwrap_or(default_flight_url);
             tracing::trace!("Connecting to SpiceAI with flight url: {url}");
+
+            verify_endpoint_connection(&url)
+                .await
+                .map_err(|e| super::Error::UnableToCreateDataConnector { source: e.into() })?;
+
             let api_key = secret.get("key").unwrap_or_default();
             let flight_client = FlightClient::new(url.as_str(), "", api_key)
                 .await