refactor: add origin information to Column

sqlx-postgres/src/column.rs

@@ -2,6 +2,7 @@ use crate::ext::ustr::UStr;
 use crate::{PgTypeInfo, Postgres};
 
 pub(crate) use sqlx_core::column::{Column, ColumnIndex};
+use sqlx_core::column::ColumnOrigin;
 
 #[derive(Debug, Clone)]
 #[cfg_attr(feature = "offline", derive(serde::Serialize, serde::Deserialize))]
@@ -9,6 +10,10 @@ pub struct PgColumn {
     pub(crate) ordinal: usize,
     pub(crate) name: UStr,
     pub(crate) type_info: PgTypeInfo,
+
+    #[cfg_attr(feature = "offline", serde(default))]
+    pub(crate) origin: ColumnOrigin,
+    
     #[cfg_attr(feature = "offline", serde(skip))]
     pub(crate) relation_id: Option<crate::types::Oid>,
     #[cfg_attr(feature = "offline", serde(skip))]
@@ -51,4 +56,8 @@ impl Column for PgColumn {
     fn type_info(&self) -> &PgTypeInfo {
         &self.type_info
     }
+
+    fn origin(&self) -> ColumnOrigin {
+        self.origin.clone()
+    }
 }

sqlx-postgres/src/connection/describe.rs

@@ -1,3 +1,4 @@
+use std::collections::btree_map;
 use crate::error::Error;
 use crate::ext::ustr::UStr;
 use crate::io::StatementId;
@@ -13,6 +14,9 @@ use crate::{PgColumn, PgConnection, PgTypeInfo};
 use smallvec::SmallVec;
 use sqlx_core::query_builder::QueryBuilder;
 use std::sync::Arc;
+use sqlx_core::column::{ColumnOrigin, TableColumn};
+use sqlx_core::hash_map;
+use crate::connection::TableColumns;
 
 /// Describes the type of the `pg_type.typtype` column
 ///
@@ -121,6 +125,12 @@ impl PgConnection {
             let type_info = self
                 .maybe_fetch_type_info_by_oid(field.data_type_id, should_fetch)
                 .await?;
+            
+            let origin = if let (Some(relation_oid), Some(attribute_no)) = (field.relation_id, field.relation_attribute_no) {
+                self.maybe_fetch_column_origin(relation_oid, attribute_no, should_fetch).await?
+            } else {
+                ColumnOrigin::Expression
+            };
 
             let column = PgColumn {
                 ordinal: index,
@@ -128,6 +138,7 @@ impl PgConnection {
                 type_info,
                 relation_id: field.relation_id,
                 relation_attribute_no: field.relation_attribute_no,
+                origin,
             };
 
             columns.push(column);
@@ -189,6 +200,54 @@ impl PgConnection {
             Ok(PgTypeInfo(PgType::DeclareWithOid(oid)))
         }
     }
+    
+    async fn maybe_fetch_column_origin(
+        &mut self, 
+        relation_id: Oid, 
+        attribute_no: i16,
+        should_fetch: bool,
+    ) -> Result<ColumnOrigin, Error> {
+        let mut table_columns = match self.cache_table_to_column_names.entry(relation_id) {
+            hash_map::Entry::Occupied(table_columns) => {
+                table_columns.into_mut()
+            },
+            hash_map::Entry::Vacant(vacant) => {
+                if !should_fetch { return Ok(ColumnOrigin::Unknown); }
+                
+                let table_name: String = query_scalar("SELECT $1::oid::regclass::text")
+                    .bind(relation_id)
+                    .fetch_one(&mut *self)
+                    .await?;
+                
+                vacant.insert(TableColumns {
+                    table_name: table_name.into(),
+                    columns: Default::default(),
+                })
+            }
+        };
+        
+        let column_name = match table_columns.columns.entry(attribute_no) {
+            btree_map::Entry::Occupied(occupied) => Arc::clone(occupied.get()),
+            btree_map::Entry::Vacant(vacant) => {
+                if !should_fetch { return Ok(ColumnOrigin::Unknown); }
+                
+                let column_name: String = query_scalar(
+                    "SELECT attname FROM pg_attribute WHERE attrelid = $1 AND attnum = $2"
+                )
+                    .bind(relation_id)
+                    .bind(attribute_no)
+                    .fetch_one(&mut *self)
+                    .await?;
+                
+                Arc::clone(vacant.insert(column_name.into()))
+            }
+        };
+        
+        Ok(ColumnOrigin::Table(TableColumn {
+            table: table_columns.table_name.clone(),
+            name: column_name
+        }))
+    }
 
     async fn fetch_type_by_oid(&mut self, oid: Oid) -> Result<PgTypeInfo, Error> {
         let (name, typ_type, category, relation_id, element, base_type): (

sqlx-postgres/src/connection/mod.rs

@@ -1,3 +1,4 @@
+use std::collections::BTreeMap;
 use std::fmt::{self, Debug, Formatter};
 use std::sync::Arc;
 
@@ -61,6 +62,7 @@ pub struct PgConnectionInner {
     cache_type_info: HashMap<Oid, PgTypeInfo>,
     cache_type_oid: HashMap<UStr, Oid>,
     cache_elem_type_to_array: HashMap<Oid, Oid>,
+    cache_table_to_column_names: HashMap<Oid, TableColumns>,
 
     // number of ReadyForQuery messages that we are currently expecting
     pub(crate) pending_ready_for_query_count: usize,
@@ -72,6 +74,12 @@ pub struct PgConnectionInner {
     log_settings: LogSettings,
 }
 
+pub(crate) struct TableColumns {
+    table_name: Arc<str>,
+    /// Attribute number -> name.
+    columns: BTreeMap<i16, Arc<str>>,
+}
+
 impl PgConnection {
     /// the version number of the server in `libpq` format
     pub fn server_version_num(&self) -> Option<u32> {