fix(sqlalchemy-spanner): group multi-column unique constraints and add get_multi_unique_constraints

packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py

@@ -1669,40 +1669,102 @@ def get_table_names(self, connection, schema=None, **kw):
         return table_names
 
     @engine_to_connection
-    def get_unique_constraints(self, connection, table_name, schema=None, **kw):
-        """Get the table unique constraints.
+    def get_multi_unique_constraints(
+        self, connection, schema=None, filter_names=None, scope=None, kind=None, **kw
+    ):
+        """
+        Return information about unique constraints in all tables
+        in the given schema.
 
         The method is used by SQLAlchemy introspection systems.
 
         Args:
             connection (sqlalchemy.engine.base.Connection):
                 SQLAlchemy connection or engine object.
-            table_name (str): Name of the table to introspect.
             schema (str): Optional. Schema name
+            filter_names (Sequence[str]): Optional. Optionally return information
+                only for the objects listed here.
+            scope (sqlalchemy.engine.reflection.ObjectScope): Optional. Specifies
+                if constraints of default, temporary or any tables
+                should be reflected. Spanner does not support temporary.
+            kind (sqlalchemy.engine.reflection.ObjectKind): Optional. Specifies the
+                type of objects to reflect.
 
         Returns:
-            dict: Dict with the unique constraints' descriptions.
+            dictionary: a dictionary where the keys are two-tuple schema,table-name
+                and the values are list of dictionaries, each representing
+                a unique constraint definition. The columns of each constraint
+                are listed in the order of the constraint.
+                The schema is ``None`` if no schema is provided.
         """
+        table_filter_query = self._get_table_filter_query(filter_names, "tc", True)
+        schema_filter_query = " tc.TABLE_SCHEMA = '{schema}' AND ".format(
+            schema=schema or ""
+        )
+        table_type_query = self._get_table_type_query(kind, True)
+
         sql = """
-SELECT ccu.CONSTRAINT_NAME, ccu.COLUMN_NAME
-FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS tc
-JOIN INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE AS ccu
-    USING (TABLE_CATALOG, TABLE_SCHEMA, CONSTRAINT_NAME)
-WHERE
-    tc.TABLE_NAME="{table_name}"
-    AND tc.TABLE_SCHEMA="{table_schema}"
-    AND tc.CONSTRAINT_TYPE = "UNIQUE"
-    AND tc.CONSTRAINT_NAME IS NOT NULL
-""".format(table_schema=schema or "", table_name=table_name)
+            SELECT
+                tc.TABLE_SCHEMA,
+                tc.TABLE_NAME,
+                tc.CONSTRAINT_NAME,
+                ARRAY_AGG(kcu.COLUMN_NAME ORDER BY kcu.ORDINAL_POSITION)
+            FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS tc
+            JOIN INFORMATION_SCHEMA.KEY_COLUMN_USAGE AS kcu
+                ON  kcu.CONSTRAINT_CATALOG = tc.CONSTRAINT_CATALOG
+                AND kcu.CONSTRAINT_SCHEMA = tc.CONSTRAINT_SCHEMA
+                AND kcu.CONSTRAINT_NAME = tc.CONSTRAINT_NAME
+            JOIN INFORMATION_SCHEMA.TABLES AS t
+                ON  t.TABLE_CATALOG = tc.TABLE_CATALOG
+                AND t.TABLE_SCHEMA = tc.TABLE_SCHEMA
+                AND t.TABLE_NAME = tc.TABLE_NAME
+            WHERE
+                {table_filter_query}
+                {table_type_query}
+                {schema_filter_query}
+                tc.CONSTRAINT_TYPE = "UNIQUE"
+                AND tc.CONSTRAINT_NAME IS NOT NULL
+            GROUP BY tc.TABLE_SCHEMA, tc.TABLE_NAME, tc.CONSTRAINT_NAME
+        """.format(
+            table_filter_query=table_filter_query,
+            table_type_query=table_type_query,
+            schema_filter_query=schema_filter_query,
+        )
 
-        cols = []
         with connection.connection.database.snapshot() as snap:
-            rows = snap.execute_sql(sql)
+            rows = list(snap.execute_sql(sql))
+            result_dict = {}
 
             for row in rows:
-                cols.append({"name": row[0], "column_names": [row[1]]})
+                table_info = result_dict.setdefault((row[0] or None, row[1]), [])
+                # The emulator lists the backing unique constraint of a foreign
+                # key once per referencing foreign key, which repeats its columns.
+                column_names = list(dict.fromkeys(row[3]))
+                table_info.append({"name": row[2], "column_names": column_names})
+
+        return result_dict
+
+    @engine_to_connection
+    def get_unique_constraints(self, connection, table_name, schema=None, **kw):
+        """Get the table unique constraints.
 
-        return cols
+        The method is used by SQLAlchemy introspection systems.
+
+        Args:
+            connection (sqlalchemy.engine.base.Connection):
+                SQLAlchemy connection or engine object.
+            table_name (str): Name of the table to introspect.
+            schema (str): Optional. Schema name
+
+        Returns:
+            list: Dicts, each of which describes a unique constraint.
+        """
+        kind = None if not USING_SQLACLCHEMY_20 else ObjectKind.ANY
+        constraints = self.get_multi_unique_constraints(
+            connection, schema=schema, filter_names=[table_name], kind=kind
+        )
+        schema = schema or None
+        return constraints.get((schema, table_name), [])
 
     @engine_to_connection
     def has_table(self, connection, table_name, schema=None, **kw):

packages/sqlalchemy-spanner/tests/test_suite_20.py

@@ -1137,11 +1137,60 @@ def test_get_multi_columns(
                 result.pop((schema, t), None)
             self._check_table_dict(result, exp, self._required_column_keys)
 
-    @pytest.mark.skip(
-        "Requires an introspection method to be implemented in SQLAlchemy first"
-    )
-    def test_get_multi_unique_constraints(self):
-        pass
+    def test_get_multi_unique_constraints(self, metadata, connection):
+        """
+        SPANNER OVERRIDE:
+
+        Spanner doesn't support unique constraints, so the tables of the
+        original test have none. Instead, a foreign key that refers to columns
+        other than the primary key gets a backing unique constraint from Spanner.
+        Checks that it is reflected as one constraint with all its columns, also
+        when more than one foreign key refers to the same columns.
+        """
+        Table(
+            "uq_parent",
+            metadata,
+            Column("id", sqlalchemy.Integer, primary_key=True),
+            Column("a", sqlalchemy.String(10)),
+            Column("b", sqlalchemy.String(10)),
+        )
+        for child in ("uq_child", "uq_child2"):
+            Table(
+                child,
+                metadata,
+                Column("id", sqlalchemy.Integer, primary_key=True),
+                Column("x", sqlalchemy.String(10)),
+                Column("y", sqlalchemy.String(10)),
+                sqlalchemy.ForeignKeyConstraint(
+                    ["x", "y"], ["uq_parent.a", "uq_parent.b"]
+                ),
+            )
+        metadata.create_all(connection)
+        connection.connection.commit()
+
+        insp = inspect(connection)
+        constraints = insp.get_multi_unique_constraints(
+            filter_names=["uq_parent", "uq_child", "uq_child2"]
+        )
+        eq_(constraints.get((None, "uq_child"), []), [])
+        eq_(constraints.get((None, "uq_child2"), []), [])
+        eq_(len(constraints[(None, "uq_parent")]), 1)
+        constraint = constraints[(None, "uq_parent")][0]
+        eq_(constraint["column_names"], ["a", "b"])
+
+        insp.clear_cache()
+        eq_(insp.get_unique_constraints("uq_parent"), [constraint])
+        eq_(insp.get_unique_constraints("uq_child"), [])
+
+        reflected = Table("uq_parent", MetaData(), autoload_with=connection)
+        eq_(
+            [
+                [col.name for col in uq.columns]
+                for uq in reflected.constraints
+                if isinstance(uq, sqlalchemy.UniqueConstraint)
+            ],
+            [["a", "b"]],
+        )
 
     @pytest.mark.skip(
         "Requires an introspection method to be implemented in SQLAlchemy first"

packages/sqlalchemy-spanner/tests/unit/test_dialect.py

@@ -87,6 +87,88 @@ def test_get_multi_indexes_handles_null_column_orderings_array(self):
         index_info = res[("public", "my_table")][0]
         eq_(index_info["column_sorting"], {})
 
+    def test_get_multi_unique_constraints_groups_columns(self):
+        """Test get_multi_unique_constraints groups the columns of each constraint."""
+        dialect = SpannerDialect()
+        connection = MagicMock()
+        mock_snapshot = MagicMock()
+        # Mock rows: schema, table, constraint_name, column_names. The emulator
+        # repeats the columns of the backing unique constraint of a foreign key
+        # for each referencing foreign key.
+        mock_snapshot.execute_sql.return_value = [
+            ["", "parent", "IDX_parent_a_b_U", ["a", "a", "b", "b"]],
+            ["", "parent", "IDX_parent_c_U", ["c"]],
+            ["my_schema", "other", "IDX_other_x_U", ["x"]],
+        ]
+        connection.connection.database.snapshot.return_value.__enter__.return_value = (
+            mock_snapshot
+        )
+
+        res = dialect.get_multi_unique_constraints(connection)
+
+        eq_(
+            res,
+            {
+                (None, "parent"): [
+                    {"name": "IDX_parent_a_b_U", "column_names": ["a", "b"]},
+                    {"name": "IDX_parent_c_U", "column_names": ["c"]},
+                ],
+                ("my_schema", "other"): [
+                    {"name": "IDX_other_x_U", "column_names": ["x"]}
+                ],
+            },
+        )
+
+    def test_get_multi_unique_constraints_sql(self):
+        """Test get_multi_unique_constraints reads all tables in one query."""
+        dialect = SpannerDialect()
+        connection = MagicMock()
+        mock_snapshot = MagicMock()
+        mock_snapshot.execute_sql.return_value = []
+        connection.connection.database.snapshot.return_value.__enter__.return_value = (
+            mock_snapshot
+        )
+
+        dialect.get_multi_unique_constraints(
+            connection, filter_names=["parent", "other"]
+        )
+
+        mock_snapshot.execute_sql.assert_called_once()
+        executed_sql = mock_snapshot.execute_sql.call_args[0][0]
+        assert "INFORMATION_SCHEMA.KEY_COLUMN_USAGE AS kcu" in executed_sql
+        assert (
+            "ARRAY_AGG(kcu.COLUMN_NAME ORDER BY kcu.ORDINAL_POSITION)" in executed_sql
+        )
+        assert (
+            "GROUP BY tc.TABLE_SCHEMA, tc.TABLE_NAME, tc.CONSTRAINT_NAME"
+            in executed_sql
+        )
+        assert "tc.table_name = 'parent' OR tc.table_name = 'other'" in executed_sql
+        assert 'tc.CONSTRAINT_TYPE = "UNIQUE"' in executed_sql
+        assert "tc.CONSTRAINT_NAME IS NOT NULL" in executed_sql
+
+    def test_get_unique_constraints(self):
+        """Test get_unique_constraints returns the constraints of one table."""
+        dialect = SpannerDialect()
+        connection = MagicMock()
+        mock_snapshot = MagicMock()
+        mock_snapshot.execute_sql.return_value = [
+            ["", "parent", "IDX_parent_a_b_U", ["a", "b"]],
+        ]
+        connection.connection.database.snapshot.return_value.__enter__.return_value = (
+            mock_snapshot
+        )
+
+        eq_(
+            dialect.get_unique_constraints(connection, "parent"),
+            [{"name": "IDX_parent_a_b_U", "column_names": ["a", "b"]}],
+        )
+        executed_sql = mock_snapshot.execute_sql.call_args[0][0]
+        assert "tc.table_name = 'parent'" in executed_sql
+
+        mock_snapshot.execute_sql.return_value = []
+        eq_(dialect.get_unique_constraints(connection, "no_constraints"), [])
+
     def test_max_size_exported(self):
         """Test MAX_SIZE export and int_from_size helper behavior."""
         from google.cloud.sqlalchemy_spanner import MAX_SIZE