From 2c4b400143302c14237aa41ffc94f81c10880bb9 Mon Sep 17 00:00:00 2001
From: "Hugh A. Miles II" <hughmil3s@gmail.com>
Date: Mon, 23 May 2022 17:58:15 -0400
Subject: [PATCH] feat: Add Certified filter to Datasets (#20136)

---
 .../src/components/ListView/types.ts          |  1 +
 .../views/CRUD/data/dataset/DatasetList.tsx   | 13 ++++++++
 superset/datasets/api.py                      |  8 +++--
 superset/datasets/filters.py                  | 18 +++++++++++
 tests/integration_tests/datasets/api_tests.py | 31 +++++++++++++++++++
 5 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/superset-frontend/src/components/ListView/types.ts b/superset-frontend/src/components/ListView/types.ts
index 53710b84d271d..f8bff90f0ee95 100644
--- a/superset-frontend/src/components/ListView/types.ts
+++ b/superset-frontend/src/components/ListView/types.ts
@@ -113,4 +113,5 @@ export enum FilterOperator {
   chartIsFav = 'chart_is_favorite',
   chartIsCertified = 'chart_is_certified',
   dashboardIsCertified = 'dashboard_is_certified',
+  datasetIsCertified = 'dataset_is_certified',
 }
diff --git a/superset-frontend/src/views/CRUD/data/dataset/DatasetList.tsx b/superset-frontend/src/views/CRUD/data/dataset/DatasetList.tsx
index 0617c48602dd3..bdc689122d05c 100644
--- a/superset-frontend/src/views/CRUD/data/dataset/DatasetList.tsx
+++ b/superset-frontend/src/views/CRUD/data/dataset/DatasetList.tsx
@@ -258,6 +258,7 @@ const DatasetList: FunctionComponent<DatasetListProps> = ({
         accessor: 'kind_icon',
         disableSortBy: true,
         size: 'xs',
+        id: 'id',
       },
       {
         Cell: ({
@@ -506,6 +507,18 @@ const DatasetList: FunctionComponent<DatasetListProps> = ({
           { label: 'Physical', value: true },
         ],
       },
+      {
+        Header: t('Certified'),
+        id: 'id',
+        urlDisplay: 'certified',
+        input: 'select',
+        operator: FilterOperator.datasetIsCertified,
+        unfilteredLabel: t('Any'),
+        selects: [
+          { label: t('Yes'), value: true },
+          { label: t('No'), value: false },
+        ],
+      },
       {
         Header: t('Search'),
         id: 'table_name',
diff --git a/superset/datasets/api.py b/superset/datasets/api.py
index 268337c090fee..fb01b6ee8c9cc 100644
--- a/superset/datasets/api.py
+++ b/superset/datasets/api.py
@@ -52,7 +52,7 @@
 from superset.datasets.commands.refresh import RefreshDatasetCommand
 from superset.datasets.commands.update import UpdateDatasetCommand
 from superset.datasets.dao import DatasetDAO
-from superset.datasets.filters import DatasetIsNullOrEmptyFilter
+from superset.datasets.filters import DatasetCertifiedFilter, DatasetIsNullOrEmptyFilter
 from superset.datasets.schemas import (
     DatasetPostSchema,
     DatasetPutSchema,
@@ -195,7 +195,11 @@ class DatasetRestApi(BaseSupersetModelRestApi):
         "owners": RelatedFieldFilter("first_name", FilterRelatedOwners),
         "database": "database_name",
     }
-    search_filters = {"sql": [DatasetIsNullOrEmptyFilter]}
+    search_filters = {
+        "sql": [DatasetIsNullOrEmptyFilter],
+        "id": [DatasetCertifiedFilter],
+    }
+    search_columns = ["id", "database", "owners", "sql", "table_name"]
     filter_rel_fields = {"database": [["id", DatabaseFilter, lambda: []]]}
     allowed_rel_fields = {"database", "owners"}
     allowed_distinct_fields = {"schema"}
diff --git a/superset/datasets/filters.py b/superset/datasets/filters.py
index 4bbe80fd4e4f1..e72eb281818ae 100644
--- a/superset/datasets/filters.py
+++ b/superset/datasets/filters.py
@@ -33,3 +33,21 @@ def apply(self, query: Query, value: bool) -> Query:
             filter_clause = not_(filter_clause)
 
         return query.filter(filter_clause)
+
+
+class DatasetCertifiedFilter(BaseFilter):  # pylint: disable=too-few-public-methods
+    name = _("Is certified")
+    arg_name = "dataset_is_certified"
+
+    def apply(self, query: Query, value: bool) -> Query:
+        check_value = '%"certification":%'
+        if value is True:
+            return query.filter(SqlaTable.extra.ilike(check_value))
+        if value is False:
+            return query.filter(
+                or_(
+                    SqlaTable.extra.notlike(check_value),
+                    SqlaTable.extra.is_(None),
+                )
+            )
+        return query
diff --git a/tests/integration_tests/datasets/api_tests.py b/tests/integration_tests/datasets/api_tests.py
index 87d0da3cad827..781ae929b743c 100644
--- a/tests/integration_tests/datasets/api_tests.py
+++ b/tests/integration_tests/datasets/api_tests.py
@@ -128,6 +128,7 @@ def create_datasets(self):
             main_db = get_main_database()
             for tables_name in self.fixture_tables_names:
                 datasets.append(self.insert_dataset(tables_name, [admin.id], main_db))
+
             yield datasets
 
             # rollback changes
@@ -1811,3 +1812,33 @@ def test_import_dataset_invalid_v0_validation(self):
                 }
             ]
         }
+
+    @pytest.mark.usefixtures("create_datasets")
+    def test_get_datasets_is_certified_filter(self):
+        """
+        Dataset API: Test custom dataset_is_certified filter
+        """
+        table_w_certification = SqlaTable(
+            table_name="foo",
+            schema=None,
+            owners=[],
+            database=get_main_database(),
+            sql=None,
+            extra='{"certification": 1}',
+        )
+        db.session.add(table_w_certification)
+        db.session.commit()
+
+        arguments = {
+            "filters": [{"col": "id", "opr": "dataset_is_certified", "value": True}]
+        }
+        self.login(username="admin")
+        uri = f"api/v1/dataset/?q={prison.dumps(arguments)}"
+        rv = self.client.get(uri)
+
+        assert rv.status_code == 200
+        response = json.loads(rv.data.decode("utf-8"))
+        assert response.get("count") == 1
+
+        db.session.delete(table_w_certification)
+        db.session.commit()