mirror of
https://github.com/neondatabase/neon.git
synced 2026-01-14 17:02:56 +00:00
6768a71c86
## Problem For #11992 I realised we need to get the type info before executing the query. This is important to know how to decode rows with custom types, eg the following query: ```sql CREATE TYPE foo AS ENUM ('foo','bar','baz'); SELECT ARRAY['foo'::foo, 'bar'::foo, 'baz'::foo] AS data; ``` Getting that to work was harder that it seems. The original tokio-postgres setup has a split between `Client` and `Connection`, where messages are passed between. Because multiple clients were supported, each client message included a dedicated response channel. Each request would be terminated by the `ReadyForQuery` message. The flow I opted to use for parsing types early would not trigger a `ReadyForQuery`. The flow is as follows: ``` PARSE "" // parse the user provided query DESCRIBE "" // describe the query, returning param/result type oids FLUSH // force postgres to flush the responses early // wait for descriptions // check if we know the types, if we don't then // setup the typeinfo query and execute it against each OID: PARSE typeinfo // prepare our typeinfo query DESCRIBE typeinfo FLUSH // force postgres to flush the responses early // wait for typeinfo statement // for each OID we don't know: BIND typeinfo EXECUTE FLUSH // wait for type info, might reveal more OIDs to inspect // close the typeinfo query, we cache the OID->type map and this is kinder to pgbouncer. CLOSE typeinfo // finally once we know all the OIDs: BIND "" // bind the user provided query - already parsed - to the user provided params EXECUTE // run the user provided query SYNC // commit the transaction ``` ## Summary of changes Please review commit by commit. The main challenge was allowing one query to issue multiple sub-queries. To do this I first made sure that the client could fully own the connection, which required removing any shared client state. I then had to replace the way responses are sent to the client, by using only a single permanent channel. This required some additional effort to track which query is being processed. Lastly I had to modify the query/typeinfo functions to not issue `sync` commands, so it would fit into the desired flow above. To note: the flow above does force an extra roundtrip into each query. I don't know yet if this has a measurable latency overhead.
91 lines
2.2 KiB
Rust
91 lines
2.2 KiB
Rust
use std::io;
|
|
|
|
use bytes::{Bytes, BytesMut};
|
|
use fallible_iterator::FallibleIterator;
|
|
use postgres_protocol2::message::backend;
|
|
use tokio_util::codec::{Decoder, Encoder};
|
|
|
|
pub enum FrontendMessage {
|
|
Raw(Bytes),
|
|
}
|
|
|
|
pub enum BackendMessage {
|
|
Normal { messages: BackendMessages },
|
|
Async(backend::Message),
|
|
}
|
|
|
|
pub struct BackendMessages(BytesMut);
|
|
|
|
impl BackendMessages {
|
|
pub fn empty() -> BackendMessages {
|
|
BackendMessages(BytesMut::new())
|
|
}
|
|
}
|
|
|
|
impl FallibleIterator for BackendMessages {
|
|
type Item = backend::Message;
|
|
type Error = io::Error;
|
|
|
|
fn next(&mut self) -> io::Result<Option<backend::Message>> {
|
|
backend::Message::parse(&mut self.0)
|
|
}
|
|
}
|
|
|
|
pub struct PostgresCodec;
|
|
|
|
impl Encoder<FrontendMessage> for PostgresCodec {
|
|
type Error = io::Error;
|
|
|
|
fn encode(&mut self, item: FrontendMessage, dst: &mut BytesMut) -> io::Result<()> {
|
|
match item {
|
|
FrontendMessage::Raw(buf) => dst.extend_from_slice(&buf),
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl Decoder for PostgresCodec {
|
|
type Item = BackendMessage;
|
|
type Error = io::Error;
|
|
|
|
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<BackendMessage>, io::Error> {
|
|
let mut idx = 0;
|
|
|
|
while let Some(header) = backend::Header::parse(&src[idx..])? {
|
|
let len = header.len() as usize + 1;
|
|
if src[idx..].len() < len {
|
|
break;
|
|
}
|
|
|
|
match header.tag() {
|
|
backend::NOTICE_RESPONSE_TAG
|
|
| backend::NOTIFICATION_RESPONSE_TAG
|
|
| backend::PARAMETER_STATUS_TAG => {
|
|
if idx == 0 {
|
|
let message = backend::Message::parse(src)?.unwrap();
|
|
return Ok(Some(BackendMessage::Async(message)));
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
|
|
idx += len;
|
|
|
|
if header.tag() == backend::READY_FOR_QUERY_TAG {
|
|
break;
|
|
}
|
|
}
|
|
|
|
if idx == 0 {
|
|
Ok(None)
|
|
} else {
|
|
Ok(Some(BackendMessage::Normal {
|
|
messages: BackendMessages(src.split_to(idx)),
|
|
}))
|
|
}
|
|
}
|
|
}
|