Include schema in filter for indexes (#17480)

Author: sethsamuel

postgres/changelog.d/17480.fixed

@@ -0,0 +1 @@
+Fixed a bug where schemas with tables of the same name were incorrectly reporting indexes of those tables multiple times

postgres/datadog_checks/postgres/metadata.py

@@ -104,9 +104,10 @@
 PG_INDEXES_QUERY = """
 SELECT indexname AS NAME,
        indexdef  AS definition,
+       schemaname,
        tablename
 FROM   pg_indexes
-WHERE  {table_names_like};
+WHERE  schemaname='{schema_name}' AND ({table_names_like});
 """
 
 PG_CHECK_FOR_FOREIGN_KEY = """
@@ -144,7 +145,7 @@
 SELECT relname,
        pg_get_partkeydef(oid) AS partition_key
 FROM   pg_class
-WHERE  relname in ({table_names});
+WHERE  oid in ({table_ids});
 """
 
 NUM_PARTITIONS_QUERY = """
@@ -295,7 +296,7 @@ def report_postgres_metadata(self):
                             tables_buffer = []
 
                             for tables in table_chunks:
-                                table_info = self._query_table_information(cursor, tables)
+                                table_info = self._query_table_information(cursor, schema['name'], tables)
 
                                 tables_buffer = [*tables_buffer, *table_info]
                                 for t in table_info:
@@ -455,7 +456,7 @@ def _query_tables_for_schema(
         return table_payloads
 
     def _query_table_information(
-        self, cursor: psycopg2.extensions.cursor, table_info: List[Dict[str, Union[str, bool]]]
+        self, cursor: psycopg2.extensions.cursor, schema_name: str, table_info: List[Dict[str, Union[str, bool]]]
     ) -> List[Dict[str, Union[str, Dict]]]:
         """
         Collect table information . Returns a dictionary
@@ -481,11 +482,10 @@ def _query_table_information(
         tables = {t.get("name"): {**t, "num_partitions": 0} for t in table_info}
         table_name_lookup = {t.get("id"): t.get("name") for t in table_info}
         table_ids = ",".join(["'{}'".format(t.get("id")) for t in table_info])
-        table_names = ",".join(["'{}'".format(t.get("name")) for t in table_info])
         table_names_like = " OR ".join(["tablename LIKE '{}%'".format(t.get("name")) for t in table_info])
 
         # Get indexes
-        cursor.execute(PG_INDEXES_QUERY.format(table_names_like=table_names_like))
+        cursor.execute(PG_INDEXES_QUERY.format(schema_name=schema_name, table_names_like=table_names_like))
         rows = cursor.fetchall()
         for row in rows:
             # Partition indexes in some versions of Postgres have appended digits for each partition
@@ -497,7 +497,7 @@ def _query_table_information(
 
         # Get partitions
         if VersionUtils.transform_version(str(self._check.version))["version.major"] != "9":
-            cursor.execute(PARTITION_KEY_QUERY.format(table_names=table_names))
+            cursor.execute(PARTITION_KEY_QUERY.format(table_ids=table_ids))
             rows = cursor.fetchall()
             for row in rows:
                 tables.get(row.get("relname"))["partition_key"] = row.get("partition_key")

postgres/tests/compose/resources/03_load_data.sh

@@ -28,6 +28,8 @@ psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" datadog_test <<-EOSQL
     SELECT * FROM persons;
     SELECT * FROM persons;
     SELECT * FROM persons;
+    CREATE SCHEMA public2;
+    CREATE TABLE public2.cities (city VARCHAR(255), country VARCHAR(255), PRIMARY KEY(city));
     GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA public TO bob;
     GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA public TO blocking_bob;
 EOSQL

postgres/tests/test_metadata.py

@@ -97,7 +97,7 @@ def test_collect_schemas(integration_check, dbm_instance, aggregator):
         # there should only two schemas, 'public' and 'datadog'. datadog is empty
         schema = database_metadata[0]['schemas'][0]
         schema_name = schema['name']
-        assert schema_name in ['public', 'datadog']
+        assert schema_name in ['public', 'public2', 'datadog']
         if schema_name == 'public':
             for table in schema['tables']:
                 tables_got.append(table['name'])
@@ -120,6 +120,7 @@ def test_collect_schemas(integration_check, dbm_instance, aggregator):
                 if table['name'] == "cities":
                     keys = list(table.keys())
                     assert_fields(keys, ["indexes", "columns", "toast_table", "id", "name", "owner"])
+                    assert len(table['indexes']) == 1
                     assert_fields(list(table['indexes'][0].keys()), ['name', 'definition'])
                 if float(POSTGRES_VERSION) >= 11:
                     if table['name'] in ('test_part', 'test_part_no_activity'):