Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

postgres: use Oid type everywhere instead of u32 #1602

Merged
merged 10 commits into from Apr 14, 2022
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