Reflect collation in types on PostgreSQL

Added support for reflection of collation in types for PostgreSQL.
The ``collation`` will be set only if different from the default
one for the type.

References #6511
Closes: #12510
Pull-request: https://github.com/sqlalchemy/sqlalchemy/pull/12510
Pull-request-sha: c0a390314eb8ff2a4f729babf25cb052b6268d0a

Change-Id: I269a194d526a0689a4b38f10456d28539c73cffb

diff --git a/doc/build/changelog/unreleased_20/6511.rst b/doc/build/changelog/unreleased_20/6511.rst
new file mode 100644 (file)
index 0000000..c23ddef
--- /dev/null
+++ b/doc/build/changelog/unreleased_20/6511.rst
@@ -0,0 +1,8 @@
+.. change::
+    :tags: usecase, postgresql
+    :tickets: 6511
+
+    Added support for reflection of collation in types for PostgreSQL.
+    The ``collation`` will be set only if different from the default
+    one for the type.
+    Pull request courtesy Denis Laxalde.

diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index 11cdcd5f94f1d2caa7eb405e47e93be213eca761..b3932ea90f490ef782c1ea2fd9b3871fd6014feb 100644 (file)
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -3799,6 +3799,32 @@ class PGDialect(default.DefaultDialect):
             .scalar_subquery()
             .label("default")
         )
+
+        # get the name of the collate when it's different from the default one
+        collate = sql.case(
+            (
+                sql.and_(
+                    pg_catalog.pg_attribute.c.attcollation != 0,
+                    select(pg_catalog.pg_type.c.typcollation)
+                    .where(
+                        pg_catalog.pg_type.c.oid
+                        == pg_catalog.pg_attribute.c.atttypid,
+                    )
+                    .correlate(pg_catalog.pg_attribute)
+                    .scalar_subquery()
+                    != pg_catalog.pg_attribute.c.attcollation,
+                ),
+                select(pg_catalog.pg_collation.c.collname)
+                .where(
+                    pg_catalog.pg_collation.c.oid
+                    == pg_catalog.pg_attribute.c.attcollation
+                )
+                .correlate(pg_catalog.pg_attribute)
+                .scalar_subquery(),
+            ),
+            else_=sql.null(),
+        ).label("collation")
+
         relkinds = self._kind_to_relkinds(kind)
         query = (
             select(
@@ -3813,6 +3839,7 @@ class PGDialect(default.DefaultDialect):
                 pg_catalog.pg_description.c.description.label("comment"),
                 generated,
                 identity,
+                collate,
             )
             .select_from(pg_catalog.pg_class)
             # NOTE: postgresql support table with no user column, meaning
@@ -3891,6 +3918,7 @@ class PGDialect(default.DefaultDialect):
         domains: Dict[str, ReflectedDomain],
         enums: Dict[str, ReflectedEnum],
         type_description: str,
+        collation: Optional[str],
     ) -> sqltypes.TypeEngine[Any]:
         """
         Attempts to reconstruct a column type defined in ischema_names based
@@ -3991,6 +4019,7 @@ class PGDialect(default.DefaultDialect):
                     domains,
                     enums,
                     type_description="DOMAIN '%s'" % domain["name"],
+                    collation=domain["collation"],
                 )
                 args = (domain["name"], data_type)
 
@@ -4023,6 +4052,9 @@ class PGDialect(default.DefaultDialect):
             )
             return sqltypes.NULLTYPE
 
+        if collation is not None:
+            kwargs["collation"] = collation
+
         data_type = schema_type(*args, **kwargs)
         if array_dim >= 1:
             # postgres does not preserve dimensionality or size of array types.
@@ -4041,11 +4073,14 @@ class PGDialect(default.DefaultDialect):
                 continue
             table_cols = columns[(schema, row_dict["table_name"])]
 
+            collation = row_dict["collation"]
+
             coltype = self._reflect_type(
                 row_dict["format_type"],
                 domains,
                 enums,
                 type_description="column '%s'" % row_dict["name"],
+                collation=collation,
             )
 
             default = row_dict["default"]

diff --git a/lib/sqlalchemy/testing/requirements.py b/lib/sqlalchemy/testing/requirements.py
index d4828712a66d8e7ef6c47f3812084b2726b1e6b1..d22e37a2a5875215e7471bdd1f532e0bed5b6eb9 100644 (file)
--- a/lib/sqlalchemy/testing/requirements.py
+++ b/lib/sqlalchemy/testing/requirements.py
@@ -689,6 +689,15 @@ class SuiteRequirements(Requirements):
         and their reflection"""
         return exclusions.closed()
 
+    @property
+    def column_collation_reflection(self):
+        """Indicates if the database support column collation reflection.
+
+        This requirement also uses ``get_order_by_collation`` to get
+        an available collation.
+        """
+        return exclusions.closed()
+
     @property
     def view_column_reflection(self):
         """target database must support retrieval of the columns in a view,

diff --git a/lib/sqlalchemy/testing/suite/test_reflection.py b/lib/sqlalchemy/testing/suite/test_reflection.py
index b4fc34a4767cfdd87c958bae5e71f97b8d12fbe0..7da5e0541401eccdfab85ffc5147d5cc79625caa 100644 (file)
--- a/lib/sqlalchemy/testing/suite/test_reflection.py
+++ b/lib/sqlalchemy/testing/suite/test_reflection.py
@@ -48,6 +48,7 @@ from ...sql.schema import BLANK_SCHEMA
 from ...testing import ComparesIndexes
 from ...testing import ComparesTables
 from ...testing import is_false
+from ...testing import is_none
 from ...testing import is_true
 from ...testing import mock
 
@@ -2489,6 +2490,29 @@ class ComponentReflectionTest(ComparesTables, OneConnectionTablesTest):
         c = insp.get_columns("unicode_comments")[0]
         eq_({c["name"]: c["comment"]}, {"emoji": "🐍🧙🝝🧙‍♂️🧙‍♀️"})
 
+    @testing.requires.column_collation_reflection
+    @testing.requires.order_by_collation
+    def test_column_collation_reflection(self, connection, metadata):
+        collation = testing.requires.get_order_by_collation(config)
+        Table(
+            "t",
+            metadata,
+            Column("collated", sa.String(collation=collation)),
+            Column("not_collated", sa.String()),
+        )
+        metadata.create_all(connection)
+
+        m2 = MetaData()
+        t2 = Table("t", m2, autoload_with=connection)
+
+        eq_(t2.c.collated.type.collation, collation)
+        is_none(t2.c.not_collated.type.collation)
+
+        insp = inspect(connection)
+        collated, not_collated = insp.get_columns("t")
+        eq_(collated["type"].collation, collation)
+        is_none(not_collated["type"].collation)
+
 
 class TableNoColumnsTest(fixtures.TestBase):
     __requires__ = ("reflect_tables_no_columns",)

diff --git a/test/dialect/postgresql/test_reflection.py b/test/dialect/postgresql/test_reflection.py
index 93a0f699cf7a71bbdce889f2f03e29a9b0d31579..534c31a860dd40a638ee522139f43716a2fd5187 100644 (file)
--- a/test/dialect/postgresql/test_reflection.py
+++ b/test/dialect/postgresql/test_reflection.py
@@ -2714,9 +2714,10 @@ class ReflectionTest(
 
 class CustomTypeReflectionTest(fixtures.TestBase):
     class CustomType:
-        def __init__(self, arg1=None, arg2=None):
+        def __init__(self, arg1=None, arg2=None, collation=None):
             self.arg1 = arg1
             self.arg2 = arg2
+            self.collation = collation
 
     ischema_names = None
 
@@ -2742,6 +2743,7 @@ class CustomTypeReflectionTest(fixtures.TestBase):
                 "format_type": sch,
                 "default": None,
                 "not_null": False,
+                "collation": "cc" if sch == "my_custom_type()" else None,
                 "comment": None,
                 "generated": "",
                 "identity_options": None,
@@ -2756,6 +2758,10 @@ class CustomTypeReflectionTest(fixtures.TestBase):
             assert isinstance(column_info["type"], self.CustomType)
             eq_(column_info["type"].arg1, args[0])
             eq_(column_info["type"].arg2, args[1])
+            if sch == "my_custom_type()":
+                eq_(column_info["type"].collation, "cc")
+            else:
+                eq_(column_info["type"].collation, None)
 
     def test_clslevel(self):
         postgresql.PGDialect.ischema_names["my_custom_type"] = self.CustomType
@@ -2784,6 +2790,7 @@ class CustomTypeReflectionTest(fixtures.TestBase):
                 "format_type": None,
                 "default": None,
                 "not_null": False,
+                "collation": None,
                 "comment": None,
                 "generated": "",
                 "identity_options": None,

diff --git a/test/requirements.py b/test/requirements.py
index cf1c0f382530c3311852e175e68dfe35dff8fef6..3b5ea9e43dbb433eb95984f45286e2276e444ec2 100644 (file)
--- a/test/requirements.py
+++ b/test/requirements.py
@@ -175,6 +175,10 @@ class DefaultRequirements(SuiteRequirements):
     def constraint_comment_reflection(self):
         return only_on(["postgresql"])
 
+    @property
+    def column_collation_reflection(self):
+        return only_on(["postgresql"])
+
     @property
     def unbounded_varchar(self):
         """Target database must support VARCHAR with no length"""