From 7dbc2e61cee8841ce510fcb2d042d2091b2ceb63 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Thu, 9 Jun 2022 16:00:30 -0700
Subject: [PATCH] Add unit test

---
 superset/result_set.py              |  7 ++-
 superset/superset_typing.py         |  8 +++-
 tests/unit_tests/result_set_test.py | 67 +++++++++++++++++++++++++++++
 3 files changed, 77 insertions(+), 5 deletions(-)
 create mode 100644 tests/unit_tests/result_set_test.py

diff --git a/superset/result_set.py b/superset/result_set.py
index a1b2cdc3abd74..725bf1449cc79 100644
--- a/superset/result_set.py
+++ b/superset/result_set.py
@@ -101,10 +101,9 @@ def __init__(  # pylint: disable=too-many-locals
 
         if cursor_description:
             # get deduped list of column names
-            column_names = dedup([col[0] for col in cursor_description])
-
-            # ensure colum names are strings (see #20137)
-            column_names = [convert_to_string(name) for name in column_names]
+            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 1af04494d0c95..ae8787d1c6913 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 0000000000000..80d7ced61ecd0
--- /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()
+    )