Add Aurora role tag to PostgreSQL metrics

postgres/assets/configuration/spec.yaml

@@ -164,7 +164,9 @@ files:
       display_priority: 0
       fleet_configurable: true
       hidden: true
-      description: Tag metrics and checks with `replication_role:<master|standby>`.
+      description: |
+        Tag metrics and checks with `replication_role:<master|standby>`. Aurora instances are also tagged with
+        `aurora_role:<writer|reader>`.
       value:
         type: boolean
         example: true

postgres/changelog.d/24901.added

@@ -0,0 +1 @@
+Add an `aurora_role` tag with writer and reader values to PostgreSQL metrics emitted from Aurora instances.

postgres/datadog_checks/postgres/postgres.py

@@ -574,6 +574,12 @@ def _get_replication_role(self):
                 # value fetched for role is of <type 'bool'>
                 return "standby" if role else "master"
 
+    def _update_replication_role_tags(self, replication_role: str) -> None:
+        self.tag_manager.set_tag('replication_role', replication_role, replace=True)
+        if self.is_aurora:
+            aurora_role = 'reader' if replication_role == 'standby' else 'writer'
+            self.tag_manager.set_tag('aurora_role', aurora_role, replace=True)
+
     def _collect_wal_metrics(self):
         if self.version >= V10:
             # _collect_stats will gather wal file metrics
@@ -786,7 +792,6 @@ def _run_query_scope(self, scope, is_custom_metrics, cols, descriptors, dbname=N
             # This happens for example when trying to get replication metrics from readers in Aurora. Let's ignore it.
             log_func(e)
             self.log.debug("Disabling replication metrics")
-            self.is_aurora = False
             self.metrics_cache.replication_metrics = {}
         except psycopg.errors.UndefinedFunction as e:
             log_func(e)
@@ -1236,7 +1241,8 @@ def check(self, _):
                 self.tag_manager.set_tag('postgresql_cluster_name', self.cluster_name, replace=True)
 
             if self._config.tag_replication_role:
-                self.tag_manager.set_tag('replication_role', self._get_replication_role(), replace=True)
+                replication_role = self._get_replication_role()
+                self._update_replication_role_tags(replication_role)
 
             tags = self.tag_manager.get_tags()
             self._send_database_instance_metadata()

postgres/tests/common.py

@@ -173,6 +173,8 @@ def _get_expected_tags(
     )
     if role:
         base_tags.append(f'replication_role:{role}')
+        if check.is_aurora:
+            base_tags.append(f'aurora_role:{"reader" if role == "standby" else "writer"}')
     if with_db:
         base_tags.append(f'db:{pg_instance["dbname"]}')
     if with_host:

postgres/tests/test_pg_integration.py

@@ -350,30 +350,69 @@ def test_session_idle_in_transaction_time(aggregator, integration_check, pg_inst
     postgres_conn.close()
 
 
-def test_unsupported_replication(aggregator, integration_check, pg_instance):
+def test_feature_not_supported_preserves_aurora_role_after_promotion(aggregator, integration_check, pg_instance):
+    pg_instance['tag_replication_role'] = True
+    pg_instance['database_instance_collection_interval'] = 0
     check = integration_check(pg_instance)
+    check.is_aurora = True
+    check._get_replication_role = mock.MagicMock(side_effect=['standby', 'master'])
     unpatched_fmt = PartialFormatter()
 
     called = []
 
-    def format_with_error(value, **kwargs):
-        if 'pg_is_in_recovery' in value:
+    def format_with_error(value: str, **kwargs: str) -> str:
+        if 'pg_stat_bgwriter' in value and not called:
             called.append(True)
             raise psycopg.errors.FeatureNotSupported("Not available")
         return unpatched_fmt.format(value, **kwargs)
 
-    # This simulate an error in the fmt function, as it's a bit hard to mock psycopg
+    # Simulate an Aurora reader rejecting a metrics query.
     with mock.patch.object(fmt, 'format', passthrough=True) as mock_fmt:
         mock_fmt.side_effect = format_with_error
         check.run()
 
-    # Verify our mocking was called
+        reader_role_tags = {'replication_role:standby', 'aurora_role:reader'}
+        db_count = aggregator.metrics('postgresql.db.count')
+        assert len(db_count) == 1
+        assert {tag for tag in db_count[0].tags if tag.startswith(('replication_role:', 'aurora_role:'))} == (
+            reader_role_tags
+        )
+
+        service_checks = aggregator.service_checks(check.SERVICE_CHECK_NAME)
+        assert len(service_checks) == 1
+        assert {
+            tag for tag in service_checks[0].tags if tag.startswith(('replication_role:', 'aurora_role:'))
+        } == reader_role_tags
+
+        dbm_metadata = aggregator.get_event_platform_events("dbm-metadata")
+        database_instance = next(event for event in dbm_metadata if event['kind'] == 'database_instance')
+        assert {
+            tag for tag in database_instance['tags'] if tag.startswith(('replication_role:', 'aurora_role:'))
+        } == reader_role_tags
+
+        aggregator.reset()
+        check.run()
+
     assert called == [True]
 
-    expected_tags = _get_expected_tags(check, pg_instance)
-    check_bgw_metrics(aggregator, expected_tags)
+    writer_role_tags = {'replication_role:master', 'aurora_role:writer'}
+    db_count = aggregator.metrics('postgresql.db.count')
+    assert len(db_count) == 1
+    assert {
+        tag for tag in db_count[0].tags if tag.startswith(('replication_role:', 'aurora_role:'))
+    } == writer_role_tags
 
-    check_common_metrics(aggregator, expected_tags=expected_tags)
+    service_checks = aggregator.service_checks(check.SERVICE_CHECK_NAME)
+    assert len(service_checks) == 1
+    assert {
+        tag for tag in service_checks[0].tags if tag.startswith(('replication_role:', 'aurora_role:'))
+    } == writer_role_tags
+
+    dbm_metadata = aggregator.get_event_platform_events("dbm-metadata")
+    database_instance = next(event for event in dbm_metadata if event['kind'] == 'database_instance')
+    assert {
+        tag for tag in database_instance['tags'] if tag.startswith(('replication_role:', 'aurora_role:'))
+    } == writer_role_tags
 
 
 def test_can_connect_service_check(aggregator, integration_check, pg_instance):