postgres: use Oid type everywhere instead of u32 (#1602)

* 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

Author: paolobarbolini

sqlx-core/src/postgres/arguments.rs

@@ -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(())

sqlx-core/src/postgres/connection/describe.rs

@@ -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};
@@ -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
@@ -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)
@@ -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
@@ -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
@@ -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 {
@@ -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
@@ -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
                 ",
@@ -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() {
@@ -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 {

sqlx-core/src/postgres/connection/establish.rs

@@ -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
@@ -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(),

sqlx-core/src/postgres/connection/executor.rs

@@ -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,
@@ -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
@@ -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());
         }

sqlx-core/src/postgres/connection/mod.rs

@@ -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;
 
@@ -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,

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> {
@@ -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);

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).
     ///

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 {

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';
@@ -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 {
@@ -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);
 
@@ -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);
 

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”.
@@ -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,
     };
 

sqlx-core/src/postgres/message/parameter_description.rs

@@ -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 {
@@ -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 })
@@ -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]