Skip to content

Commit

Permalink
postgres: use Oid type everywhere instead of u32 (#1602)
Browse files Browse the repository at this point in the history
* postgres: use Oid type instead of u32

* Make serde happy

* Expose the inner u32

* docs

* Try to fix tests

* Fix unit tests

* Fix order

* Not sure what happened here
  • Loading branch information
paolobarbolini committed Apr 14, 2022
1 parent a68957b commit a97208c
Show file tree
Hide file tree
Showing 22 changed files with 270 additions and 201 deletions.
2 changes: 1 addition & 1 deletion sqlx-core/src/postgres/arguments.rs
Expand Up @@ -96,7 +96,7 @@ impl PgArguments {

for (offset, name) in type_holes {
let oid = conn.fetch_type_id_by_name(&*name).await?;
buffer[*offset..(*offset + 4)].copy_from_slice(&oid.to_be_bytes());
buffer[*offset..(*offset + 4)].copy_from_slice(&oid.0.to_be_bytes());
}

Ok(())
Expand Down
36 changes: 20 additions & 16 deletions sqlx-core/src/postgres/connection/describe.rs
Expand Up @@ -3,6 +3,7 @@ use crate::ext::ustr::UStr;
use crate::postgres::message::{ParameterDescription, RowDescription};
use crate::postgres::statement::PgStatementMetadata;
use crate::postgres::type_info::{PgCustomType, PgType, PgTypeKind};
use crate::postgres::types::Oid;
use crate::postgres::{PgArguments, PgColumn, PgConnection, PgTypeInfo};
use crate::query_as::query_as;
use crate::query_scalar::{query_scalar, query_scalar_with};
Expand Down Expand Up @@ -147,7 +148,7 @@ impl PgConnection {

async fn maybe_fetch_type_info_by_oid(
&mut self,
oid: u32,
oid: Oid,
should_fetch: bool,
) -> Result<PgTypeInfo, Error> {
// first we check if this is a built-in type
Expand Down Expand Up @@ -183,9 +184,9 @@ impl PgConnection {
}
}

fn fetch_type_by_oid(&mut self, oid: u32) -> BoxFuture<'_, Result<PgTypeInfo, Error>> {
fn fetch_type_by_oid(&mut self, oid: Oid) -> BoxFuture<'_, Result<PgTypeInfo, Error>> {
Box::pin(async move {
let (name, typ_type, category, relation_id, element, base_type): (String, i8, i8, u32, u32, u32) = query_as(
let (name, typ_type, category, relation_id, element, base_type): (String, i8, i8, Oid, Oid, Oid) = query_as(
"SELECT typname, typtype, typcategory, typrelid, typelem, typbasetype FROM pg_catalog.pg_type WHERE oid = $1",
)
.bind(oid)
Expand Down Expand Up @@ -237,7 +238,7 @@ impl PgConnection {
})
}

async fn fetch_enum_by_oid(&mut self, oid: u32, name: String) -> Result<PgTypeInfo, Error> {
async fn fetch_enum_by_oid(&mut self, oid: Oid, name: String) -> Result<PgTypeInfo, Error> {
let variants: Vec<String> = query_scalar(
r#"
SELECT enumlabel
Expand All @@ -259,12 +260,12 @@ ORDER BY enumsortorder

fn fetch_composite_by_oid(
&mut self,
oid: u32,
relation_id: u32,
oid: Oid,
relation_id: Oid,
name: String,
) -> BoxFuture<'_, Result<PgTypeInfo, Error>> {
Box::pin(async move {
let raw_fields: Vec<(String, u32)> = query_as(
let raw_fields: Vec<(String, Oid)> = query_as(
r#"
SELECT attname, atttypid
FROM pg_catalog.pg_attribute
Expand Down Expand Up @@ -296,8 +297,8 @@ ORDER BY attnum

fn fetch_domain_by_oid(
&mut self,
oid: u32,
base_type: u32,
oid: Oid,
base_type: Oid,
name: String,
) -> BoxFuture<'_, Result<PgTypeInfo, Error>> {
Box::pin(async move {
Expand All @@ -313,11 +314,11 @@ ORDER BY attnum

fn fetch_range_by_oid(
&mut self,
oid: u32,
oid: Oid,
name: String,
) -> BoxFuture<'_, Result<PgTypeInfo, Error>> {
Box::pin(async move {
let element_oid: u32 = query_scalar(
let element_oid: Oid = query_scalar(
r#"
SELECT rngsubtype
FROM pg_catalog.pg_range
Expand All @@ -338,13 +339,13 @@ WHERE rngtypid = $1
})
}

pub(crate) async fn fetch_type_id_by_name(&mut self, name: &str) -> Result<u32, Error> {
pub(crate) async fn fetch_type_id_by_name(&mut self, name: &str) -> Result<Oid, Error> {
if let Some(oid) = self.cache_type_oid.get(name) {
return Ok(*oid);
}

// language=SQL
let (oid,): (u32,) = query_as(
let (oid,): (Oid,) = query_as(
"
SELECT oid FROM pg_catalog.pg_type WHERE typname ILIKE $1
",
Expand All @@ -362,7 +363,7 @@ SELECT oid FROM pg_catalog.pg_type WHERE typname ILIKE $1

pub(crate) async fn get_nullable_for_columns(
&mut self,
stmt_id: u32,
stmt_id: Oid,
meta: &PgStatementMetadata,
) -> Result<Vec<Option<bool>>, Error> {
if meta.columns.is_empty() {
Expand Down Expand Up @@ -424,10 +425,13 @@ SELECT oid FROM pg_catalog.pg_type WHERE typname ILIKE $1
/// and returns `None` for all others.
async fn nullables_from_explain(
&mut self,
stmt_id: u32,
stmt_id: Oid,
params_len: usize,
) -> Result<Vec<Option<bool>>, Error> {
let mut explain = format!("EXPLAIN (VERBOSE, FORMAT JSON) EXECUTE sqlx_s_{}", stmt_id);
let mut explain = format!(
"EXPLAIN (VERBOSE, FORMAT JSON) EXECUTE sqlx_s_{}",
stmt_id.0
);
let mut comma = false;

if params_len > 0 {
Expand Down
3 changes: 2 additions & 1 deletion sqlx-core/src/postgres/connection/establish.rs
Expand Up @@ -7,6 +7,7 @@ use crate::postgres::connection::{sasl, stream::PgStream, tls};
use crate::postgres::message::{
Authentication, BackendKeyData, MessageFormat, Password, ReadyForQuery, Startup,
};
use crate::postgres::types::Oid;
use crate::postgres::{PgConnectOptions, PgConnection};

// https://www.postgresql.org/docs/current/protocol-flow.html#id-1.10.5.7.3
Expand Down Expand Up @@ -143,7 +144,7 @@ impl PgConnection {
transaction_status,
transaction_depth: 0,
pending_ready_for_query_count: 0,
next_statement_id: 1,
next_statement_id: Oid(1),
cache_statement: StatementCache::new(options.statement_cache_capacity),
cache_type_oid: HashMap::new(),
cache_type_info: HashMap::new(),
Expand Down
7 changes: 4 additions & 3 deletions sqlx-core/src/postgres/connection/executor.rs
Expand Up @@ -8,6 +8,7 @@ use crate::postgres::message::{
};
use crate::postgres::statement::PgStatementMetadata;
use crate::postgres::type_info::PgType;
use crate::postgres::types::Oid;
use crate::postgres::{
statement::PgStatement, PgArguments, PgConnection, PgQueryResult, PgRow, PgTypeInfo,
PgValueFormat, Postgres,
Expand All @@ -24,9 +25,9 @@ async fn prepare(
sql: &str,
parameters: &[PgTypeInfo],
metadata: Option<Arc<PgStatementMetadata>>,
) -> Result<(u32, Arc<PgStatementMetadata>), Error> {
) -> Result<(Oid, Arc<PgStatementMetadata>), Error> {
let id = conn.next_statement_id;
conn.next_statement_id = conn.next_statement_id.wrapping_add(1);
conn.next_statement_id.incr_one();

// build a list of type OIDs to send to the database in the PARSE command
// we have not yet started the query sequence, so we are *safe* to cleanly make
Expand Down Expand Up @@ -169,7 +170,7 @@ impl PgConnection {
// optional metadata that was provided by the user, this means they are reusing
// a statement object
metadata: Option<Arc<PgStatementMetadata>>,
) -> Result<(u32, Arc<PgStatementMetadata>), Error> {
) -> Result<(Oid, Arc<PgStatementMetadata>), Error> {
if let Some(statement) = self.cache_statement.get_mut(sql) {
return Ok((*statement).clone());
}
Expand Down
9 changes: 5 additions & 4 deletions sqlx-core/src/postgres/connection/mod.rs
Expand Up @@ -15,6 +15,7 @@ use crate::postgres::message::{
Close, Message, MessageFormat, Query, ReadyForQuery, Terminate, TransactionStatus,
};
use crate::postgres::statement::PgStatementMetadata;
use crate::postgres::types::Oid;
use crate::postgres::{PgConnectOptions, PgTypeInfo, Postgres};
use crate::transaction::Transaction;

Expand Down Expand Up @@ -46,14 +47,14 @@ pub struct PgConnection {

// sequence of statement IDs for use in preparing statements
// in PostgreSQL, the statement is prepared to a user-supplied identifier
next_statement_id: u32,
next_statement_id: Oid,

// cache statement by query string to the id and columns
cache_statement: StatementCache<(u32, Arc<PgStatementMetadata>)>,
cache_statement: StatementCache<(Oid, Arc<PgStatementMetadata>)>,

// cache user-defined types by id <-> info
cache_type_info: HashMap<u32, PgTypeInfo>,
cache_type_oid: HashMap<UStr, u32>,
cache_type_info: HashMap<Oid, PgTypeInfo>,
cache_type_oid: HashMap<UStr, Oid>,

// number of ReadyForQuery messages that we are currently expecting
pub(crate) pending_ready_for_query_count: usize,
Expand Down
14 changes: 8 additions & 6 deletions sqlx-core/src/postgres/io/buf_mut.rs
@@ -1,11 +1,13 @@
use crate::postgres::types::Oid;

pub trait PgBufMutExt {
fn put_length_prefixed<F>(&mut self, f: F)
where
F: FnOnce(&mut Vec<u8>);

fn put_statement_name(&mut self, id: u32);
fn put_statement_name(&mut self, id: Oid);

fn put_portal_name(&mut self, id: Option<u32>);
fn put_portal_name(&mut self, id: Option<Oid>);
}

impl PgBufMutExt for Vec<u8> {
Expand All @@ -29,22 +31,22 @@ impl PgBufMutExt for Vec<u8> {

// writes a statement name by ID
#[inline]
fn put_statement_name(&mut self, id: u32) {
fn put_statement_name(&mut self, id: Oid) {
// N.B. if you change this don't forget to update it in ../describe.rs
self.extend(b"sqlx_s_");

self.extend(itoa::Buffer::new().format(id).as_bytes());
self.extend(itoa::Buffer::new().format(id.0).as_bytes());

self.push(0);
}

// writes a portal name by ID
#[inline]
fn put_portal_name(&mut self, id: Option<u32>) {
fn put_portal_name(&mut self, id: Option<Oid>) {
if let Some(id) = id {
self.extend(b"sqlx_p_");

self.extend(itoa::Buffer::new().format(id).as_bytes());
self.extend(itoa::Buffer::new().format(id.0).as_bytes());
}

self.push(0);
Expand Down
5 changes: 3 additions & 2 deletions sqlx-core/src/postgres/message/bind.rs
@@ -1,14 +1,15 @@
use crate::io::Encode;
use crate::postgres::io::PgBufMutExt;
use crate::postgres::types::Oid;
use crate::postgres::PgValueFormat;

#[derive(Debug)]
pub struct Bind<'a> {
/// The ID of the destination portal (`None` selects the unnamed portal).
pub portal: Option<u32>,
pub portal: Option<Oid>,

/// The id of the source prepared statement.
pub statement: u32,
pub statement: Oid,

/// The parameter format codes. Each must presently be zero (text) or one (binary).
///
Expand Down
5 changes: 3 additions & 2 deletions sqlx-core/src/postgres/message/close.rs
@@ -1,14 +1,15 @@
use crate::io::Encode;
use crate::postgres::io::PgBufMutExt;
use crate::postgres::types::Oid;

const CLOSE_PORTAL: u8 = b'P';
const CLOSE_STATEMENT: u8 = b'S';

#[derive(Debug)]
#[allow(dead_code)]
pub enum Close {
Statement(u32),
Portal(u32),
Statement(Oid),
Portal(Oid),
}

impl Encode<'_> for Close {
Expand Down
9 changes: 5 additions & 4 deletions sqlx-core/src/postgres/message/describe.rs
@@ -1,5 +1,6 @@
use crate::io::Encode;
use crate::postgres::io::PgBufMutExt;
use crate::postgres::types::Oid;

const DESCRIBE_PORTAL: u8 = b'P';
const DESCRIBE_STATEMENT: u8 = b'S';
Expand All @@ -10,10 +11,10 @@ const DESCRIBE_STATEMENT: u8 = b'S';
#[allow(dead_code)]
pub enum Describe {
UnnamedStatement,
Statement(u32),
Statement(Oid),

UnnamedPortal,
Portal(u32),
Portal(Oid),
}

impl Encode<'_> for Describe {
Expand Down Expand Up @@ -54,7 +55,7 @@ fn test_encode_describe_portal() {
const EXPECTED: &[u8] = b"D\0\0\0\x0EPsqlx_p_5\0";

let mut buf = Vec::new();
let m = Describe::Portal(5);
let m = Describe::Portal(Oid(5));

m.encode(&mut buf);

Expand All @@ -78,7 +79,7 @@ fn test_encode_describe_statement() {
const EXPECTED: &[u8] = b"D\0\0\0\x0ESsqlx_s_5\0";

let mut buf = Vec::new();
let m = Describe::Statement(5);
let m = Describe::Statement(Oid(5));

m.encode(&mut buf);

Expand Down
5 changes: 3 additions & 2 deletions sqlx-core/src/postgres/message/execute.rs
@@ -1,9 +1,10 @@
use crate::io::Encode;
use crate::postgres::io::PgBufMutExt;
use crate::postgres::types::Oid;

pub struct Execute {
/// The id of the portal to execute (`None` selects the unnamed portal).
pub portal: Option<u32>,
pub portal: Option<Oid>,

/// Maximum number of rows to return, if portal contains a query
/// that returns rows (ignored otherwise). Zero denotes “no limit”.
Expand All @@ -28,7 +29,7 @@ fn test_encode_execute() {

let mut buf = Vec::new();
let m = Execute {
portal: Some(5),
portal: Some(Oid(5)),
limit: 2,
};

Expand Down
9 changes: 5 additions & 4 deletions sqlx-core/src/postgres/message/parameter_description.rs
Expand Up @@ -3,10 +3,11 @@ use smallvec::SmallVec;

use crate::error::Error;
use crate::io::Decode;
use crate::postgres::types::Oid;

#[derive(Debug)]
pub struct ParameterDescription {
pub types: SmallVec<[u32; 6]>,
pub types: SmallVec<[Oid; 6]>,
}

impl Decode<'_> for ParameterDescription {
Expand All @@ -15,7 +16,7 @@ impl Decode<'_> for ParameterDescription {
let mut types = SmallVec::with_capacity(cnt as usize);

for _ in 0..cnt {
types.push(buf.get_u32());
types.push(Oid(buf.get_u32()));
}

Ok(Self { types })
Expand All @@ -29,8 +30,8 @@ fn test_decode_parameter_description() {
let m = ParameterDescription::decode(DATA.into()).unwrap();

assert_eq!(m.types.len(), 2);
assert_eq!(m.types[0], 0x0000_0000);
assert_eq!(m.types[1], 0x0000_0500);
assert_eq!(m.types[0], Oid(0x0000_0000));
assert_eq!(m.types[1], Oid(0x0000_0500));
}

#[test]
Expand Down

0 comments on commit a97208c

Please sign in to comment.