[superset] branch master updated: fix: ensure column name in description is string (#20340)

[mi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

michaelsmolina pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git


The following commit(s) were added to refs/heads/master by this push:
     new f3b289d3c3 fix: ensure column name in description is string (#20340)
f3b289d3c3 is described below

commit f3b289d3c333fe2351e9fbac6fa85b875cb1897c
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Tue Jun 21 06:10:46 2022 -0700

    fix: ensure column name in description is string (#20340)
    
    * fix: ensure column name in description is string
    
    * Add unit test
---
 superset/result_set.py              | 17 +++++++++-
 superset/superset_typing.py         |  8 ++++-
 tests/unit_tests/result_set_test.py | 67 +++++++++++++++++++++++++++++++++++++
 3 files changed, 90 insertions(+), 2 deletions(-)

diff --git a/superset/result_set.py b/superset/result_set.py
index 008699d683..725bf1449c 100644
--- a/superset/result_set.py
+++ b/superset/result_set.py
@@ -71,6 +71,19 @@ def destringify(obj: str) -> Any:
     return json.loads(obj)
 
 
+def convert_to_string(value: Any) -> str:
+    """
+    Used to ensure column names from the cursor description are strings.
+    """
+    if isinstance(value, str):
+        return value
+
+    if isinstance(value, bytes):
+        return value.decode("utf-8")
+
+    return str(value)
+
+
 class SupersetResultSet:
     def __init__(  # pylint: disable=too-many-locals
         self,
@@ -88,7 +101,9 @@ class SupersetResultSet:
 
         if cursor_description:
             # get deduped list of column names
-            column_names = dedup([col[0] for col in cursor_description])
+            column_names = dedup(
+                [convert_to_string(col[0]) for col in cursor_description]
+            )
 
             # fix cursor descriptor with the deduped names
             deduped_cursor_desc = [
diff --git a/superset/superset_typing.py b/superset/superset_typing.py
index 1af04494d0..ae8787d1c6 100644
--- a/superset/superset_typing.py
+++ b/superset/superset_typing.py
@@ -69,7 +69,13 @@ class ResultSetColumnType(TypedDict):
 
 CacheConfig = Dict[str, Any]
 DbapiDescriptionRow = Tuple[
-    str, str, Optional[str], Optional[str], Optional[int], Optional[int], bool
+    Union[str, bytes],
+    str,
+    Optional[str],
+    Optional[str],
+    Optional[int],
+    Optional[int],
+    bool,
 ]
 DbapiDescription = Union[List[DbapiDescriptionRow], Tuple[DbapiDescriptionRow, ...]]
 DbapiResult = Sequence[Union[List[Any], Tuple[Any, ...]]]
diff --git a/tests/unit_tests/result_set_test.py b/tests/unit_tests/result_set_test.py
new file mode 100644
index 0000000000..80d7ced61e
--- /dev/null
+++ b/tests/unit_tests/result_set_test.py
@@ -0,0 +1,67 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# pylint: disable=import-outside-toplevel, unused-argument
+
+
+def test_column_names_as_bytes(app_context: None) -> None:
+    """
+    Test that we can handle column names as bytes.
+    """
+    from superset.db_engine_specs.redshift import RedshiftEngineSpec
+    from superset.result_set import SupersetResultSet
+
+    data = (
+        [
+            "2016-01-26",
+            392.002014,
+            397.765991,
+            390.575012,
+            392.153015,
+            392.153015,
+            58147000,
+        ],
+        [
+            "2016-01-27",
+            392.444,
+            396.842987,
+            391.782013,
+            394.971985,
+            394.971985,
+            47424400,
+        ],
+    )
+    description = [
+        (b"date", 1043, None, None, None, None, None),
+        (b"open", 701, None, None, None, None, None),
+        (b"high", 701, None, None, None, None, None),
+        (b"low", 701, None, None, None, None, None),
+        (b"close", 701, None, None, None, None, None),
+        (b"adj close", 701, None, None, None, None, None),
+        (b"volume", 20, None, None, None, None, None),
+    ]
+    result_set = SupersetResultSet(data, description, RedshiftEngineSpec)  # type: ignore
+
+    assert (
+        result_set.to_pandas_df().to_markdown()
+        == """
+|    | date       |    open |    high |     low |   close |   adj close |   volume |
+|---:|:-----------|--------:|--------:|--------:|--------:|------------:|---------:|
+|  0 | 2016-01-26 | 392.002 | 397.766 | 390.575 | 392.153 |     392.153 | 58147000 |
+|  1 | 2016-01-27 | 392.444 | 396.843 | 391.782 | 394.972 |     394.972 | 47424400 |
+    """.strip()
+    )