Fix unique_constraints (#152)

Signed-off-by: Pablo Takara <[email protected]>
ibmdb · Dec 18, 2024 · bfe67ae · bfe67ae
1 parent fdb3618
commit bfe67ae
Showing 1 changed file with 59 additions and 25 deletions.
diff --git a/ibm_db_sa/reflection.py b/ibm_db_sa/reflection.py
@@ -462,24 +462,41 @@ def get_unique_constraints(self, connection, table_name, schema=None, **kw):
         table_name = self.denormalize_name(table_name)
         syskeycol = self.sys_keycoluse
         sysconst = self.sys_tabconst
-        query = sql.select(syskeycol.c.constname, syskeycol.c.colname).\
-            where(and_(
-              syskeycol.c.constname == sysconst.c.constname,
-              sysconst.c.tabname == table_name,
-              sysconst.c.tabschema == current_schema,
-              sysconst.c.type == 'U')).\
-            order_by(syskeycol.c.constname)
+        query = (
+            sql.select(syskeycol.c.constname, syskeycol.c.colname)
+            .select_from(
+                join(
+                    syskeycol,
+                    sysconst,
+                    and_(
+                        syskeycol.c.constname == sysconst.c.constname,
+                        syskeycol.c.tabschema == sysconst.c.tabschema,
+                        syskeycol.c.tabname == sysconst.c.tabname,
+                    ),
+                )
+            )
+            .where(
+                and_(
+                    sysconst.c.tabname == table_name,
+                    sysconst.c.tabschema == current_schema,
+                    sysconst.c.type == "U",
+                )
+            )
+            .order_by(syskeycol.c.constname)
+        )
         uniqueConsts = []
         currConst = None
         for r in connection.execute(query):
             if currConst == r[0]:
-                uniqueConsts[-1]['column_names'].append(self.normalize_name(r[1]))
+                uniqueConsts[-1]["column_names"].append(self.normalize_name(r[1]))
             else:
                 currConst = r[0]
-                uniqueConsts.append({
-                        'name': self.normalize_name(currConst),
-                        'column_names': [self.normalize_name(r[1])],
-                    })
+                uniqueConsts.append(
+                    {
+                        "name": self.normalize_name(currConst),
+                        "column_names": [self.normalize_name(r[1])],
+                    }
+                )
         return uniqueConsts
 
 
@@ -779,7 +796,7 @@ def get_indexes(self, connection, table_name, schema=None, **kw):
         current_schema = self.denormalize_name(
                                     schema or self.default_schema_name)
         table_name = self.denormalize_name(table_name)
-        
+
         sysidx = self.sys_indexes
         syskey = self.sys_keys
 
@@ -1170,22 +1187,39 @@ def get_unique_constraints(self, connection, table_name, schema=None, **kw):
         table_name = self.denormalize_name(table_name)
         syskeycol = self.sys_keycoluse
         sysconst = self.sys_tabconst
-        query = sql.select(syskeycol.c.constname, syskeycol.c.colname).\
-            where(and_(
-                syskeycol.c.constname == sysconst.c.constname,
-                sysconst.c.tabname == table_name,
-                sysconst.c.tabschema == current_schema,
-                sysconst.c.type == 'U')).\
-            order_by(syskeycol.c.constname)
+        query = (
+            sql.select(syskeycol.c.constname, syskeycol.c.colname)
+            .select_from(
+                join(
+                    syskeycol,
+                    sysconst,
+                    and_(
+                        syskeycol.c.constname == sysconst.c.constname,
+                        syskeycol.c.tabschema == sysconst.c.tabschema,
+                        syskeycol.c.tabname == sysconst.c.tabname,
+                    ),
+                )
+            )
+            .where(
+                and_(
+                    sysconst.c.tabname == table_name,
+                    sysconst.c.tabschema == current_schema,
+                    sysconst.c.type == "U",
+                )
+            )
+            .order_by(syskeycol.c.constname)
+        )
         uniqueConsts = []
         currConst = None
         for r in connection.execute(query):
             if currConst == r[0]:
-                uniqueConsts[-1]['column_names'].append(self.normalize_name(r[1]))
+                uniqueConsts[-1]["column_names"].append(self.normalize_name(r[1]))
             else:
                 currConst = r[0]
-                uniqueConsts.append({
-                        'name': self.normalize_name(currConst),
-                        'column_names': [self.normalize_name(r[1])],
-                    })
+                uniqueConsts.append(
+                    {
+                        "name": self.normalize_name(currConst),
+                        "column_names": [self.normalize_name(r[1])],
+                    }
+                )
         return uniqueConsts