Add data type to athena query runner (#7112)

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
This commit is contained in:
Zach Liu
2024-08-06 23:36:58 -04:00
committed by GitHub
parent b1fe2d4162
commit 285c2b6e56
2 changed files with 31 additions and 12 deletions

View File

@@ -75,7 +75,9 @@ class TestGlueSchema(TestCase):
{"DatabaseName": "test1"},
)
with self.stubber:
assert query_runner.get_schema() == [{"columns": ["row_id"], "name": "test1.jdbc_table"}]
assert query_runner.get_schema() == [
{"columns": [{"name": "row_id", "type": "int"}], "name": "test1.jdbc_table"}
]
def test_partitioned_table(self):
"""
@@ -124,7 +126,12 @@ class TestGlueSchema(TestCase):
{"DatabaseName": "test1"},
)
with self.stubber:
assert query_runner.get_schema() == [{"columns": ["sk", "category"], "name": "test1.partitioned_table"}]
assert query_runner.get_schema() == [
{
"columns": [{"name": "sk", "type": "int"}, {"name": "category", "type": "int"}],
"name": "test1.partitioned_table",
}
]
def test_view(self):
query_runner = Athena({"glue": True, "region": "mars-east-1"})
@@ -156,7 +163,7 @@ class TestGlueSchema(TestCase):
{"DatabaseName": "test1"},
)
with self.stubber:
assert query_runner.get_schema() == [{"columns": ["sk"], "name": "test1.view"}]
assert query_runner.get_schema() == [{"columns": [{"name": "sk", "type": "int"}], "name": "test1.view"}]
def test_dodgy_table_does_not_break_schema_listing(self):
"""
@@ -196,7 +203,9 @@ class TestGlueSchema(TestCase):
{"DatabaseName": "test1"},
)
with self.stubber:
assert query_runner.get_schema() == [{"columns": ["region"], "name": "test1.csv"}]
assert query_runner.get_schema() == [
{"columns": [{"name": "region", "type": "string"}], "name": "test1.csv"}
]
def test_no_storage_descriptor_table(self):
"""
@@ -312,6 +321,6 @@ class TestGlueSchema(TestCase):
)
with self.stubber:
assert query_runner.get_schema() == [
{"columns": ["row_id"], "name": "test1.jdbc_table"},
{"columns": ["row_id"], "name": "test2.jdbc_table"},
{"columns": [{"name": "row_id", "type": "int"}], "name": "test1.jdbc_table"},
{"columns": [{"name": "row_id", "type": "int"}], "name": "test2.jdbc_table"},
]