sqlness_runner/
client.rs

1// Copyright 2023 Greptime Team
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 std::net::SocketAddr;
16use std::time::Duration;
17
18use client::error::ServerSnafu;
19use client::{
20    Client, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, Database, OutputData, RecordBatches,
21};
22use common_error::ext::ErrorExt;
23use common_query::Output;
24use mysql::prelude::Queryable;
25use mysql::{Conn as MySqlClient, Row as MySqlRow};
26use tokio_postgres::{Client as PgClient, SimpleQueryMessage as PgRow};
27
28use crate::util::retry_with_backoff;
29
30/// A client that can connect to GreptimeDB using multiple protocols.
31pub struct MultiProtocolClient {
32    grpc_client: Database,
33    pg_client: PgClient,
34    mysql_client: MySqlClient,
35}
36
37/// The result of a MySQL query.
38pub enum MysqlSqlResult {
39    AffectedRows(u64),
40    Rows(Vec<MySqlRow>),
41}
42
43impl MultiProtocolClient {
44    /// Connect to the GreptimeDB server using multiple protocols.
45    ///
46    /// # Arguments
47    ///
48    /// * `grpc_server_addr` - The address of the GreptimeDB server.
49    /// * `pg_server_addr` - The address of the Postgres server.
50    /// * `mysql_server_addr` - The address of the MySQL server.
51    ///
52    /// # Returns
53    ///
54    /// A `MultiProtocolClient` instance.
55    ///
56    /// # Panics
57    ///
58    /// Panics if the server addresses are invalid or the connection fails.
59    pub async fn connect(
60        grpc_server_addr: &str,
61        pg_server_addr: &str,
62        mysql_server_addr: &str,
63    ) -> MultiProtocolClient {
64        let grpc_client = Database::new(
65            DEFAULT_CATALOG_NAME,
66            DEFAULT_SCHEMA_NAME,
67            Client::with_query_and_control_managers(
68                Default::default(),
69                Default::default(),
70                vec![grpc_server_addr],
71            ),
72        );
73        let pg_client = create_postgres_client(pg_server_addr).await;
74        let mysql_client = create_mysql_client(mysql_server_addr).await;
75        MultiProtocolClient {
76            grpc_client,
77            pg_client,
78            mysql_client,
79        }
80    }
81
82    /// Reconnect the MySQL client.
83    pub async fn reconnect_mysql_client(&mut self, mysql_server_addr: &str) {
84        self.mysql_client = create_mysql_client(mysql_server_addr).await;
85    }
86
87    /// Reconnect the Postgres client.
88    pub async fn reconnect_pg_client(&mut self, pg_server_addr: &str) {
89        self.pg_client = create_postgres_client(pg_server_addr).await;
90    }
91
92    /// Execute a query on the Postgres server.
93    pub async fn postgres_query(&mut self, query: &str) -> Result<Vec<PgRow>, String> {
94        match self.pg_client.simple_query(query).await {
95            Ok(rows) => Ok(rows),
96            Err(e) => Err(format!("Failed to execute query, encountered: {:?}", e)),
97        }
98    }
99
100    /// Execute a query on the MySQL server.
101    pub async fn mysql_query(&mut self, query: &str) -> Result<MysqlSqlResult, String> {
102        let result = self.mysql_client.query_iter(query);
103        match result {
104            Ok(result) => {
105                let mut rows = vec![];
106                let affected_rows = result.affected_rows();
107                for row in result {
108                    match row {
109                        Ok(r) => rows.push(r),
110                        Err(e) => {
111                            return Err(format!("Failed to parse query result, err: {:?}", e));
112                        }
113                    }
114                }
115
116                if rows.is_empty() {
117                    Ok(MysqlSqlResult::AffectedRows(affected_rows))
118                } else {
119                    Ok(MysqlSqlResult::Rows(rows))
120                }
121            }
122            Err(e) => Err(format!("Failed to execute query, err: {:?}", e)),
123        }
124    }
125
126    /// Execute a query on the GreptimeDB server.
127    pub async fn grpc_query(&mut self, query: &str) -> Result<Output, client::Error> {
128        let query_str = query.trim().to_lowercase();
129        if query_str.starts_with("use ") {
130            // use [db]
131            let database = query
132                .split_ascii_whitespace()
133                .nth(1)
134                .expect("Illegal `USE` statement: expecting a database.")
135                .trim_end_matches(';');
136            self.grpc_client.set_schema(database);
137            Ok(Output::new_with_affected_rows(0))
138        } else if query_str.starts_with("set time_zone")
139            || query_str.starts_with("set session time_zone")
140            || query_str.starts_with("set local time_zone")
141        {
142            // set time_zone='xxx'
143            let timezone = query
144                .split('=')
145                .nth(1)
146                .expect("Illegal `SET TIMEZONE` statement: expecting a timezone expr.")
147                .trim()
148                .strip_prefix('\'')
149                .unwrap()
150                .strip_suffix("';")
151                .unwrap();
152
153            self.grpc_client.set_timezone(timezone);
154            Ok(Output::new_with_affected_rows(0))
155        } else {
156            let mut result = self.grpc_client.sql(&query).await;
157            if let Ok(Output {
158                data: OutputData::Stream(stream),
159                ..
160            }) = result
161            {
162                match RecordBatches::try_collect(stream).await {
163                    Ok(recordbatches) => {
164                        result = Ok(Output::new_with_record_batches(recordbatches));
165                    }
166                    Err(e) => {
167                        let status_code = e.status_code();
168                        let msg = e.output_msg();
169                        result = ServerSnafu {
170                            code: status_code,
171                            msg,
172                        }
173                        .fail();
174                    }
175                }
176            }
177
178            result
179        }
180    }
181}
182
183/// Create a Postgres client with retry.
184///
185/// # Panics
186///
187/// Panics if the Postgres server address is invalid or the connection fails.
188async fn create_postgres_client(pg_server_addr: &str) -> PgClient {
189    let sockaddr: SocketAddr = pg_server_addr.parse().unwrap_or_else(|_| {
190        panic!("Failed to parse the Postgres server address {pg_server_addr}. Please check if the address is in the format of `ip:port`.")
191    });
192    let mut config = tokio_postgres::config::Config::new();
193    config.host(sockaddr.ip().to_string());
194    config.port(sockaddr.port());
195    config.dbname(DEFAULT_SCHEMA_NAME);
196
197    retry_with_backoff(
198        || async {
199            config
200                .connect(tokio_postgres::NoTls)
201                .await
202                .map(|(client, conn)| {
203                    tokio::spawn(conn);
204                    client
205                })
206        },
207        3,
208        Duration::from_millis(500),
209    )
210    .await
211    .unwrap_or_else(|_| {
212        panic!("Failed to connect to Postgres server. Please check if the server is running.")
213    })
214}
215
216/// Create a MySQL client with retry.
217///
218/// # Panics
219///
220/// Panics if the MySQL server address is invalid or the connection fails.
221async fn create_mysql_client(mysql_server_addr: &str) -> MySqlClient {
222    let sockaddr: SocketAddr = mysql_server_addr.parse().unwrap_or_else(|_| {
223        panic!("Failed to parse the MySQL server address {mysql_server_addr}. Please check if the address is in the format of `ip:port`.")
224    });
225    let ops = mysql::OptsBuilder::new()
226        .ip_or_hostname(Some(sockaddr.ip().to_string()))
227        .tcp_port(sockaddr.port())
228        .db_name(Some(DEFAULT_SCHEMA_NAME));
229
230    retry_with_backoff(
231        || async { mysql::Conn::new(ops.clone()) },
232        3,
233        Duration::from_millis(500),
234    )
235    .await
236    .unwrap_or_else(|_| {
237        panic!("Failed to connect to MySQL server. Please check if the server is running.")
238    })
239}
sqlness_runner/client.rs

sqlness_runner/
client.rs
