From 796a2a69241e7e7dda7b09f31ed2de5ca5ecef15 Mon Sep 17 00:00:00 2001
From: Maxime Beauchemin <maximebeauchemin@gmail.com>
Date: Tue, 20 Oct 2020 21:34:46 -0700
Subject: [PATCH] fix: error around latest partition in BigQuery (#11274)

* fix: error around latest partition in BigQuery

* lint

* Going with a backend-first approach

* fix test

* add an extra test
---
 .../src/SqlLab/components/TableElement.jsx    |  2 +-
 superset/db_engine_specs/base.py              | 12 +++++++++++
 superset/db_engine_specs/bigquery.py          | 19 +++++++++++++++++
 superset/models/core.py                       |  3 ++-
 tests/db_engine_specs/bigquery_tests.py       | 21 +++++++++++++++++++
 5 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/superset-frontend/src/SqlLab/components/TableElement.jsx b/superset-frontend/src/SqlLab/components/TableElement.jsx
index aeca231e7..0332aa5df 100644
--- a/superset-frontend/src/SqlLab/components/TableElement.jsx
+++ b/superset-frontend/src/SqlLab/components/TableElement.jsx
@@ -110,7 +110,7 @@ class TableElement extends React.PureComponent {
           />
         );
       }
-      let latest = Object.entries(table.partitions.latest).map(
+      let latest = Object.entries(table.partitions?.latest || []).map(
         ([key, value]) => `${key}=${value}`,
       );
       latest = latest.join('/');
diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py
index 651b98f6b..0aee290fe 100644
--- a/superset/db_engine_specs/base.py
+++ b/superset/db_engine_specs/base.py
@@ -376,6 +376,18 @@ class BaseEngineSpec:  # pylint: disable=too-many-public-methods
             return type_code.upper()
         return None
 
+    @classmethod
+    def normalize_indexes(cls, indexes: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Normalizes indexes for more consistency across db engines
+
+        noop by default
+
+        :param indexes: Raw indexes as returned by SQLAlchemy
+        :return: cleaner, more aligned index definition
+        """
+        return indexes
+
     @classmethod
     def extra_table_metadata(
         cls, database: "Database", table_name: str, schema_name: str
diff --git a/superset/db_engine_specs/bigquery.py b/superset/db_engine_specs/bigquery.py
index 1a111cd3b..7c783a2f8 100644
--- a/superset/db_engine_specs/bigquery.py
+++ b/superset/db_engine_specs/bigquery.py
@@ -129,6 +129,25 @@ class BigQueryEngineSpec(BaseEngineSpec):
         """
         return "_" + hashlib.md5(label.encode("utf-8")).hexdigest()
 
+    @classmethod
+    def normalize_indexes(cls, indexes: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Normalizes indexes for more consistency across db engines
+
+        :param indexes: Raw indexes as returned by SQLAlchemy
+        :return: cleaner, more aligned index definition
+        """
+        normalized_idxs = []
+        # Fixing a bug/behavior observed in pybigquery==0.4.15 where
+        # the index's `column_names` == [None]
+        # Here we're returning only non-None indexes
+        for ix in indexes:
+            column_names = ix.get("column_names") or []
+            ix["column_names"] = [col for col in column_names if col is not None]
+            if ix["column_names"]:
+                normalized_idxs.append(ix)
+        return normalized_idxs
+
     @classmethod
     def extra_table_metadata(
         cls, database: "Database", table_name: str, schema_name: str
diff --git a/superset/models/core.py b/superset/models/core.py
index 36bed2a1e..acc96f46f 100755
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -619,7 +619,8 @@ class Database(
     def get_indexes(
         self, table_name: str, schema: Optional[str] = None
     ) -> List[Dict[str, Any]]:
-        return self.inspector.get_indexes(table_name, schema)
+        indexes = self.inspector.get_indexes(table_name, schema)
+        return self.db_engine_spec.normalize_indexes(indexes)
 
     def get_pk_constraint(
         self, table_name: str, schema: Optional[str] = None
diff --git a/tests/db_engine_specs/bigquery_tests.py b/tests/db_engine_specs/bigquery_tests.py
index 988575dff..8e9ae7d0a 100644
--- a/tests/db_engine_specs/bigquery_tests.py
+++ b/tests/db_engine_specs/bigquery_tests.py
@@ -122,6 +122,27 @@ class TestBigQueryDbEngineSpec(TestDbEngineSpec):
         )
         self.assertEqual(result, expected_result)
 
+    def test_normalize_indexes(self):
+        """
+        DB Eng Specs (bigquery): Test extra table metadata
+        """
+        indexes = [{"name": "partition", "column_names": [None], "unique": False}]
+        normalized_idx = BigQueryEngineSpec.normalize_indexes(indexes)
+        self.assertEqual(normalized_idx, [])
+
+        indexes = [{"name": "partition", "column_names": ["dttm"], "unique": False}]
+        normalized_idx = BigQueryEngineSpec.normalize_indexes(indexes)
+        self.assertEqual(normalized_idx, indexes)
+
+        indexes = [
+            {"name": "partition", "column_names": ["dttm", None], "unique": False}
+        ]
+        normalized_idx = BigQueryEngineSpec.normalize_indexes(indexes)
+        self.assertEqual(
+            normalized_idx,
+            [{"name": "partition", "column_names": ["dttm"], "unique": False}],
+        )
+
     def test_df_to_sql(self):
         """
         DB Eng Specs (bigquery): Test DataFrame to SQL contract