risingwave_connector/source/adbc_snowflake/
schema.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use adbc_core::Statement as _;
16use anyhow::Context;
17use risingwave_common::array::arrow::arrow_schema_58 as arrow_schema;
18
19use super::AdbcSnowflakeProperties;
20use crate::error::ConnectorResult;
21
22impl AdbcSnowflakeProperties {
23    /// Get the Arrow schema from the Snowflake table.
24    /// This is used for schema inference when creating tables.
25    ///
26    /// **Important**: We use a `LIMIT 0` query instead of ADBC's `get_table_schema` API
27    /// because `get_table_schema` may return different types than the actual query results.
28    /// For example, Snowflake NUMBER columns may be reported as Int64 by `get_table_schema`
29    /// but returned as Decimal128 in actual query results. Using a real query ensures
30    /// consistency between schema inference (used by frontend) and data fetching (used by executor).
31    ///
32    /// The column order in the returned schema matches the column order in the Snowflake table.
33    pub fn get_arrow_schema(&self) -> ConnectorResult<arrow_schema::Schema> {
34        let database = self.create_database()?;
35        let mut connection = self.create_connection(&database)?;
36
37        // Use LIMIT 0 query to get the actual schema that will be returned by queries.
38        // This ensures schema inference matches actual data types returned by Snowflake.
39        let query = format!("SELECT * FROM {} LIMIT 0", self.table_ref());
40        let mut statement = self.create_statement(&mut connection, &query)?;
41
42        let reader = statement
43            .execute()
44            .context("Failed to execute schema query")?;
45
46        let schema = reader.schema();
47
48        Ok((*schema).clone())
49    }
50}
risingwave_connector/source/adbc_snowflake/schema.rs

risingwave_connector/source/adbc_snowflake/
schema.rs
