Add filtering dbs by regex capability in pgbouncer integration (DataDog#23111)

* Add filtering dbs by regex capability in pgbouncer integration

* PGBouncer - Regenerated conf.yaml.example

* PGBouncer - filter by regex, fix default value

* PGBouncer - Make new attribute field_configurable like the other attributes

* Pgbouncer - Added changelog entry

Author: yannh

pgbouncer/assets/configuration/spec.yaml

@@ -18,6 +18,18 @@ files:
         value:
           example: postgresql://<USERNAME>:<PASSWORD>@<HOSTNAME>:<PORT>/<DATABASE_URL>?sslmode=require
           type: string
+      - name: database_filter_regex
+        description: |
+          Collect metrics only for PgBouncer database names matching this regular expression.
+
+          The regex is applied to the PgBouncer database name exposed by the admin views.
+          Global metrics that are not tied to a specific database, such as configuration metrics,
+          are still collected.
+        value:
+          example: ^(datadog_test|dogs)$
+          type: string
+          display_default: null
+        fleet_configurable: true
       - name: host
         description: If `database_url` is not used, set up the host to connect to with the `host` parameter.
         value:

pgbouncer/changelog.d/23111.added

@@ -0,0 +1 @@
+Add filtering databases by regex capability.

pgbouncer/datadog_checks/pgbouncer/config_models/instance.py

@@ -36,6 +36,7 @@ class InstanceConfig(BaseModel):
     )
     collect_per_client_metrics: Optional[bool] = None
     collect_per_server_metrics: Optional[bool] = None
+    database_filter_regex: Optional[str] = None
     database_url: Optional[str] = None
     disable_generic_tags: Optional[bool] = None
     empty_default_hostname: Optional[bool] = None

pgbouncer/datadog_checks/pgbouncer/data/conf.yaml.example

@@ -23,6 +23,15 @@ instances:
     #
     # database_url: postgresql://<USERNAME>:<PASSWORD>@<HOSTNAME>:<PORT>/<DATABASE_URL>?sslmode=require
 
+    ## @param database_filter_regex - string - optional
+    ## Collect metrics only for PgBouncer database names matching this regular expression.
+    ##
+    ## The regex is applied to the PgBouncer database name exposed by the admin views.
+    ## Global metrics that are not tied to a specific database, such as configuration metrics,
+    ## are still collected.
+    #
+    # database_filter_regex: ^(datadog_test|dogs)$
+
     ## @param host - string - optional
     ## If `database_url` is not used, set up the host to connect to with the `host` parameter.
     #

pgbouncer/datadog_checks/pgbouncer/pgbouncer.py

@@ -38,9 +38,17 @@ def __init__(self, name, init_config, instances):
         self.password = self.instance.get('password', '')
         self.tags = self.instance.get('tags', [])
         self.database_url = self.instance.get('database_url')
+        self.database_filter_regex = self.instance.get('database_filter_regex', '')
         self.use_cached = is_affirmative(self.instance.get('use_cached', True))
         self.collect_per_client_metrics = is_affirmative(self.instance.get('collect_per_client_metrics', False))
         self.collect_per_server_metrics = is_affirmative(self.instance.get('collect_per_server_metrics', False))
+        self.database_filter = None
+
+        if self.database_filter_regex:
+            try:
+                self.database_filter = re.compile(self.database_filter_regex)
+            except re.error as e:
+                raise ConfigurationError('Invalid database_filter_regex: {}'.format(e))
 
         if not self.database_url:
             if not self.host:
@@ -100,6 +108,9 @@ def _collect_stats(self, db):
                         elif row.get('database') == self.DB_NAME:
                             continue
 
+                        if not self._should_collect_row(row):
+                            continue
+
                         tags = list(self.tags)
                         tags += ["%s:%s" % (tag, row[column]) for (column, tag) in descriptors if column in row]
                         for column, (name, reporter) in metrics:
@@ -136,6 +147,23 @@ def iter_rows(self, cursor):
 
             row_num += 1
 
+    def _get_row_database_name(self, row):
+        if 'name' in row:
+            return row['name']
+        if 'database' in row:
+            return row['database']
+        return None
+
+    def _should_collect_row(self, row):
+        if not self.database_filter:
+            return True
+
+        database_name = self._get_row_database_name(row)
+        if database_name is None:
+            return True
+
+        return self.database_filter.search(database_name) is not None
+
     def _get_connect_kwargs(self):
         """
         Get the params to pass to psycopg.connect() based on passed-in vals

pgbouncer/tests/test_pgbouncer_unit.py

@@ -27,6 +27,41 @@ def test_config_missing_user(instance):
         PgBouncer('pgbouncer', {}, [instance])
 
 
+@pytest.mark.unit
+def test_config_invalid_database_filter_regex(instance):
+    instance['database_filter_regex'] = '('
+
+    with pytest.raises(ConfigurationError, match='Invalid database_filter_regex'):
+        PgBouncer('pgbouncer', {}, [instance])
+
+
+@pytest.mark.unit
+def test_database_filter_matches_database_rows(instance):
+    instance['database_filter_regex'] = '^dogs$'
+    check = PgBouncer('pgbouncer', {}, [instance])
+
+    assert check._should_collect_row({'database': 'dogs'})
+    assert not check._should_collect_row({'database': 'datadog_test'})
+
+
+@pytest.mark.unit
+def test_database_filter_matches_show_databases_rows(instance):
+    instance['database_filter_regex'] = '^dogs$'
+    check = PgBouncer('pgbouncer', {}, [instance])
+
+    assert check._get_row_database_name({'name': 'dogs', 'database': 'postgres_dogs'}) == 'dogs'
+    assert check._should_collect_row({'name': 'dogs'})
+    assert not check._should_collect_row({'name': 'datadog_test'})
+
+
+@pytest.mark.unit
+def test_database_filter_keeps_global_rows(instance):
+    instance['database_filter_regex'] = '^dogs$'
+    check = PgBouncer('pgbouncer', {}, [instance])
+
+    assert check._should_collect_row({'key': 'max_client_conn', 'value': '100'})
+
+
 @pytest.mark.unit
 @pytest.mark.parametrize('use_cached', [True, False])
 def test_connection_cleanup_on_error(instance, use_cached):