fix: Handling of column types for Presto, Trino, et al. (#28653)

This commit is contained in:
John Bodley
2024-05-28 08:56:38 -07:00
committed by GitHub
parent a59bad83d4
commit 4ff17409ab
4 changed files with 85 additions and 39 deletions

View File

@@ -14,9 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=too-many-lines
from __future__ import annotations
import contextlib
@@ -118,7 +116,7 @@ def get_children(column: ResultSetColumnType) -> list[ResultSetColumnType]:
pattern = re.compile(r"(?P<type>\w+)\((?P<children>.*)\)")
if not column["type"]:
raise ValueError
match = pattern.match(column["type"])
match = pattern.match(cast(str, column["type"]))
if not match:
raise Exception( # pylint: disable=broad-exception-raised
f"Unable to parse column type {column['type']}"
@@ -538,6 +536,10 @@ class PrestoBaseEngineSpec(BaseEngineSpec, metaclass=ABCMeta):
for col_name, value in zip(col_names, values):
col_type = column_type_by_name.get(col_name)
if isinstance(col_type, str):
col_type_class = getattr(types, col_type, None)
col_type = col_type_class() if col_type_class else None
if isinstance(col_type, types.DATE):
col_type = Date()
elif isinstance(col_type, types.TIMESTAMP):