diff --git a/flightsql/src/planner.rs b/flightsql/src/planner.rs index e080abcf89..7d7448cbb3 100644 --- a/flightsql/src/planner.rs +++ b/flightsql/src/planner.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use arrow::{ - array::{as_string_array, ArrayRef, BinaryArray, GenericBinaryBuilder}, + array::{as_string_array, ArrayRef, BinaryArray, GenericBinaryBuilder, StringArray}, datatypes::{DataType, Field, Schema, SchemaRef}, error::ArrowError, ipc::writer::IpcWriteOptions, @@ -558,16 +558,22 @@ static GET_TABLES_SCHEMA_WITH_TABLE_SCHEMA: Lazy = Lazy::new(|| { }); /// Return a `LogicalPlan` for GetTableTypes -/// -/// In the future this could be made more efficient by building the -/// response directly from the IOx catalog rather than running an -/// entire DataFusion plan. async fn plan_get_table_types(ctx: &IOxSessionContext) -> Result { - let query = "SELECT DISTINCT table_type FROM information_schema.tables ORDER BY table_type"; - - Ok(ctx.sql_to_logical_plan(query).await?) + Ok(ctx.batch_to_logical_plan(TABLE_TYPES_RECORD_BATCH.clone())?) } /// The schema for GetTableTypes -static GET_TABLE_TYPE_SCHEMA: Lazy = - Lazy::new(|| Schema::new(vec![Field::new("table_type", DataType::Utf8, false)])); +static GET_TABLE_TYPE_SCHEMA: Lazy = Lazy::new(|| { + Arc::new(Schema::new(vec![Field::new( + "table_type", + DataType::Utf8, + false, + )])) +}); + +static TABLE_TYPES_RECORD_BATCH: Lazy = Lazy::new(|| { + // https://github.com/apache/arrow-datafusion/blob/26b8377b0690916deacf401097d688699026b8fb/datafusion/core/src/catalog/information_schema.rs#L285-L287 + // IOx doesn't support LOCAL TEMPORARY yet + let table_type = Arc::new(StringArray::from_iter_values(["BASE TABLE", "VIEW"])) as ArrayRef; + RecordBatch::try_new(Arc::clone(&GET_TABLE_TYPE_SCHEMA), vec![table_type]).unwrap() +}); diff --git a/influxdb_iox/tests/end_to_end_cases/flightsql.rs b/influxdb_iox/tests/end_to_end_cases/flightsql.rs index 605bf11953..1762e04190 100644 --- a/influxdb_iox/tests/end_to_end_cases/flightsql.rs +++ b/influxdb_iox/tests/end_to_end_cases/flightsql.rs @@ -488,6 +488,70 @@ async fn flightsql_get_table_types() { .await } +#[tokio::test] +async fn flightsql_get_table_types_matches_information_schema() { + test_helpers::maybe_start_logging(); + let database_url = maybe_skip_integration!(); + + let table_name = "the_table"; + + // Set up the cluster ==================================== + let mut cluster = MiniCluster::create_shared2(database_url).await; + + StepTest::new( + &mut cluster, + vec![ + Step::WriteLineProtocol(format!( + "{table_name},tag1=A,tag2=B val=42i 123456\n\ + {table_name},tag1=A,tag2=C val=43i 123457" + )), + Step::Custom(Box::new(move |state: &mut StepTestState| { + async move { + let mut client = flightsql_client(state.cluster()); + + // output of get_table_types is built manually in + // IOx, so it is important it remains in sync with + // the actual contents of the information schema + fn no_filter() -> Option { + None + } + let stream = client + .get_table_types() + .await + .unwrap(); + let get_table_types_batches = collect_stream(stream).await; + let get_table_types_output = batches_to_sorted_lines(&get_table_types_batches); + + let sql = "SELECT DISTINCT table_type FROM information_schema.tables ORDER BY table_type"; + + let stream = client.query(sql).await.unwrap(); + let information_schema_batches = collect_stream(stream).await; + let information_schema_output = + batches_to_sorted_lines(&information_schema_batches); + + insta::assert_yaml_snapshot!( + get_table_types_output, + @r###" + --- + - +------------+ + - "| table_type |" + - +------------+ + - "| BASE TABLE |" + - "| VIEW |" + - +------------+ + "### + ); + + assert_eq!(get_table_types_output, information_schema_output); + } + .boxed() + })), + ], + ) + .run() + .await +} + #[tokio::test] async fn flightsql_get_db_schemas() { test_helpers::maybe_start_logging(); diff --git a/influxdb_iox_client/src/client/flightsql.rs b/influxdb_iox_client/src/client/flightsql.rs index cf59aba68b..717c05872b 100644 --- a/influxdb_iox_client/src/client/flightsql.rs +++ b/influxdb_iox_client/src/client/flightsql.rs @@ -120,8 +120,13 @@ impl FlightSqlClient { /// Step 2: Fetch the results described in the [`FlightInfo`] /// /// This implementation does not support alternate endpoints - pub async fn query(&mut self, query: String) -> Result { - let msg = CommandStatementQuery { query }; + pub async fn query( + &mut self, + query: impl Into + Send, + ) -> Result { + let msg = CommandStatementQuery { + query: query.into(), + }; self.do_get_with_cmd(msg.as_any()).await }