risingwave_batch_executors/executor/
postgres_query.rs

1// Copyright 2024 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 anyhow::Context;
16use futures_async_stream::try_stream;
17use futures_util::stream::StreamExt;
18use risingwave_common::array::DataChunk;
19use risingwave_common::catalog::{Field, Schema};
20use risingwave_common::row::OwnedRow;
21use risingwave_common::types::{DataType, Datum, Decimal, ScalarImpl};
22use risingwave_common::util::chunk_coalesce::DataChunkBuilder;
23use risingwave_connector::connector_common::{PgConnectionConfig, create_pg_client};
24use risingwave_pb::batch_plan::plan_node::NodeBody;
25use tokio_postgres;
26
27use crate::error::BatchError;
28use crate::executor::{BoxedExecutor, BoxedExecutorBuilder, Executor, ExecutorBuilder};
29
30/// `PostgresQuery` executor. Runs a query against a Postgres database.
31pub struct PostgresQueryExecutor {
32    schema: Schema,
33    config: PgConnectionConfig,
34    query: String,
35    identity: String,
36    chunk_size: usize,
37}
38
39impl Executor for PostgresQueryExecutor {
40    fn schema(&self) -> &risingwave_common::catalog::Schema {
41        &self.schema
42    }
43
44    fn identity(&self) -> &str {
45        &self.identity
46    }
47
48    fn execute(self: Box<Self>) -> super::BoxedDataChunkStream {
49        self.do_execute().boxed()
50    }
51}
52
53pub fn postgres_row_to_owned_row(
54    row: tokio_postgres::Row,
55    schema: &Schema,
56) -> Result<OwnedRow, BatchError> {
57    let mut datums = vec![];
58    for i in 0..schema.fields.len() {
59        let rw_field = &schema.fields[i];
60        let name = rw_field.name.as_str();
61        let datum = postgres_cell_to_scalar_impl(&row, &rw_field.data_type, i, name)?;
62        datums.push(datum);
63    }
64    Ok(OwnedRow::new(datums))
65}
66
67// TODO(kwannoel): Support more types, see postgres connector's ScalarAdapter.
68fn postgres_cell_to_scalar_impl(
69    row: &tokio_postgres::Row,
70    data_type: &DataType,
71    i: usize,
72    name: &str,
73) -> Result<Datum, BatchError> {
74    let datum = match data_type {
75        DataType::Boolean
76        | DataType::Int16
77        | DataType::Int32
78        | DataType::Int64
79        | DataType::Float32
80        | DataType::Float64
81        | DataType::Date
82        | DataType::Time
83        | DataType::Timestamp
84        | DataType::Timestamptz
85        | DataType::Jsonb
86        | DataType::Interval
87        | DataType::Varchar
88        | DataType::Bytea => {
89            // ScalarAdapter is also fine. But ScalarImpl is more efficient
90            row.try_get::<_, Option<ScalarImpl>>(i)?
91        }
92        DataType::Decimal => {
93            // Decimal is more efficient than PgNumeric in ScalarAdapter
94            let val = row.try_get::<_, Option<Decimal>>(i)?;
95            val.map(ScalarImpl::from)
96        }
97        _ => {
98            tracing::warn!(name, ?data_type, "unsupported data type, set to null");
99            None
100        }
101    };
102    Ok(datum)
103}
104
105impl PostgresQueryExecutor {
106    pub fn new(
107        schema: Schema,
108        config: PgConnectionConfig,
109        query: String,
110        identity: String,
111        chunk_size: usize,
112    ) -> Self {
113        Self {
114            schema,
115            config,
116            query,
117            identity,
118            chunk_size,
119        }
120    }
121
122    #[try_stream(ok = DataChunk, error = BatchError)]
123    async fn do_execute(self: Box<Self>) {
124        tracing::debug!("postgres_query_executor: started");
125
126        let client = create_pg_client(&self.config, None).await?;
127
128        let params: &[&str] = &[];
129        let row_stream = client
130            .query_raw(&self.query, params)
131            .await
132            .context("postgres_query received error from remote server")?;
133        let mut builder = DataChunkBuilder::new(self.schema.data_types(), self.chunk_size);
134        tracing::debug!("postgres_query_executor: query executed, start deserializing rows");
135        // deserialize the rows
136        #[for_await]
137        for row in row_stream {
138            let row = row?;
139            let owned_row = postgres_row_to_owned_row(row, &self.schema)?;
140            if let Some(chunk) = builder.append_one_row(owned_row) {
141                yield chunk;
142            }
143        }
144        if let Some(chunk) = builder.consume_all() {
145            yield chunk;
146        }
147        return Ok(());
148    }
149}
150
151pub struct PostgresQueryExecutorBuilder {}
152
153impl BoxedExecutorBuilder for PostgresQueryExecutorBuilder {
154    async fn new_boxed_executor(
155        source: &ExecutorBuilder<'_>,
156        _inputs: Vec<BoxedExecutor>,
157    ) -> crate::error::Result<BoxedExecutor> {
158        let postgres_query_node = try_match_expand!(
159            source.plan_node().get_node_body().unwrap(),
160            NodeBody::PostgresQuery
161        )?;
162
163        let port = postgres_query_node
164            .port
165            .parse::<u16>()
166            .with_context(|| format!("invalid postgres port `{}`", postgres_query_node.port))?;
167        Ok(Box::new(PostgresQueryExecutor::new(
168            Schema::from_iter(postgres_query_node.columns.iter().map(Field::from)),
169            PgConnectionConfig {
170                host: postgres_query_node.hostname.clone(),
171                port,
172                user: postgres_query_node.username.clone(),
173                password: postgres_query_node.password.clone(),
174                database: postgres_query_node.database.clone(),
175                ssl_mode: postgres_query_node.ssl_mode.parse().unwrap_or_default(),
176                ssl_root_cert: if postgres_query_node.ssl_root_cert.is_empty() {
177                    None
178                } else {
179                    Some(postgres_query_node.ssl_root_cert.clone())
180                },
181            },
182            postgres_query_node.query.clone(),
183            source.plan_node().get_identity().clone(),
184            source.context().get_config().developer.chunk_size,
185        )))
186    }
187}
risingwave_batch_executors/executor/postgres_query.rs

risingwave_batch_executors/executor/
postgres_query.rs
