Support schema in Postgres type names, generate PgHasArrayType for enums

remind101 · Jul 20, 2023 · d4ddcea · d4ddcea
1 parent c70cfaf
commit d4ddcea
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 8 deletions.
diff --git a/sqlx-macros-core/src/derives/attributes.rs b/sqlx-macros-core/src/derives/attributes.rs
@@ -236,12 +236,6 @@ pub fn check_enum_attributes(input: &DeriveInput) -> syn::Result<SqlxContainerAt
         input
     );
 
-    assert_attribute!(
-        !attributes.no_pg_array,
-        "unused #[sqlx(no_pg_array)]; derive does not emit `PgHasArrayType` impls for enums",
-        input
-    );
-
     Ok(attributes)
 }
 

diff --git a/sqlx-macros-core/src/derives/type.rs b/sqlx-macros-core/src/derives/type.rs
@@ -184,6 +184,18 @@ fn expand_derive_has_sql_type_strong_enum(
                 }
             }
         ));
+
+        if !attributes.no_pg_array {
+            let array_ty_name = pg_array_type_name(ident, attributes.type_name.as_ref());
+            tts.extend(quote!(
+                    #[automatically_derived]
+                    impl ::sqlx::postgres::PgHasArrayType for #ident {
+                        fn array_type_info() -> ::sqlx::postgres::PgTypeInfo {
+                            ::sqlx::postgres::PgTypeInfo::with_name(#array_ty_name)
+                        }
+                    }
+            ));
+        }
     }
 
     if cfg!(feature = "sqlite") {
@@ -235,3 +247,21 @@ fn type_name(ident: &Ident, explicit_name: Option<&TypeName>) -> TokenStream {
         quote_spanned!(ident.span()=> #s)
     })
 }
+
+fn pg_array_type_name(ident: &Ident, explicit_name: Option<&TypeName>) -> TokenStream {
+    explicit_name
+        .map(|tn| {
+            let s = pg_array_type_name_from_type_name(&tn.val);
+            quote! { #s }
+        })
+        .unwrap_or_else(|| {
+            let s = pg_array_type_name_from_type_name(&ident.to_string());
+            quote_spanned!(ident.span()=> #s)
+        })
+}
+
+fn pg_array_type_name_from_type_name(name: &str) -> String {
+    name.split_once('.')
+        .map(|(schema, name)| format!("{}._{}", schema, name))
+        .unwrap_or_else(|| format!("_{}", name))
+}
diff --git a/sqlx-postgres/src/connection/describe.rs b/sqlx-postgres/src/connection/describe.rs
@@ -186,7 +186,12 @@ impl PgConnection {
     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, Oid, Oid, Oid) = query_as(
-                "SELECT typname, typtype, typcategory, typrelid, typelem, typbasetype FROM pg_catalog.pg_type WHERE oid = $1",
+                r#"
+SELECT nspname || '.' || typname, typtype, typcategory, typrelid, typelem, typbasetype
+FROM pg_catalog.pg_type
+JOIN pg_catalog.pg_namespace on pg_namespace.oid = pg_type.typnamespace
+WHERE pg_type.oid = $1
+                "#,
             )
             .bind(oid)
             .fetch_one(&mut *self)

diff --git a/sqlx-postgres/src/type_info.rs b/sqlx-postgres/src/type_info.rs
@@ -1152,7 +1152,13 @@ impl PartialEq<PgType> for PgType {
             true
         } else {
             // Otherwise, perform a match on the name
-            self.name().eq_ignore_ascii_case(other.name())
+            // If either name has no schema qualification, match only on the name part.
+            // If both names have a schema qualification, match both the schema and name parts.
+            std::iter::zip(
+                self.name().rsplitn(2, '.'),
+                other.name().rsplitn(2, '.'),
+            )
+                .all(|(part, other_part)| part.eq_ignore_ascii_case(other_part))
         }
     }
 }