fix: save columns reference from sqllab save datasets flow (#24248)

This commit is contained in:
Hugh A. Miles II
2023-06-20 13:54:19 -04:00
committed by GitHub
parent fdef9cbc96
commit 93e1db4bd9
37 changed files with 489 additions and 202 deletions

View File

@@ -85,7 +85,7 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
results = PrestoEngineSpec.get_columns(inspector, "", "")
self.assertEqual(len(expected_results), len(results))
for expected_result, result in zip(expected_results, results):
self.assertEqual(expected_result[0], result["name"])
self.assertEqual(expected_result[0], result["column_name"])
self.assertEqual(expected_result[1], str(result["type"]))
def test_presto_get_column(self):
@@ -175,21 +175,21 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
def test_presto_get_fields(self):
cols = [
{"name": "column"},
{"name": "column.nested_obj"},
{"name": 'column."quoted.nested obj"'},
{"column_name": "column"},
{"column_name": "column.nested_obj"},
{"column_name": 'column."quoted.nested obj"'},
]
actual_results = PrestoEngineSpec._get_fields(cols)
expected_results = [
{"name": '"column"', "label": "column"},
{"name": '"column"."nested_obj"', "label": "column.nested_obj"},
{"column_name": '"column"', "label": "column"},
{"column_name": '"column"."nested_obj"', "label": "column.nested_obj"},
{
"name": '"column"."quoted.nested obj"',
"column_name": '"column"."quoted.nested obj"',
"label": 'column."quoted.nested obj"',
},
]
for actual_result, expected_result in zip(actual_results, expected_results):
self.assertEqual(actual_result.element.name, expected_result["name"])
self.assertEqual(actual_result.element.name, expected_result["column_name"])
self.assertEqual(actual_result.name, expected_result["label"])
@mock.patch.dict(
@@ -199,8 +199,18 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
)
def test_presto_expand_data_with_simple_structural_columns(self):
cols = [
{"name": "row_column", "type": "ROW(NESTED_OBJ VARCHAR)", "is_dttm": False},
{"name": "array_column", "type": "ARRAY(BIGINT)", "is_dttm": False},
{
"column_name": "row_column",
"name": "row_column",
"type": "ROW(NESTED_OBJ VARCHAR)",
"is_dttm": False,
},
{
"column_name": "array_column",
"name": "array_column",
"type": "ARRAY(BIGINT)",
"is_dttm": False,
},
]
data = [
{"row_column": ["a"], "array_column": [1, 2, 3]},
@@ -210,9 +220,24 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
cols, data
)
expected_cols = [
{"name": "row_column", "type": "ROW(NESTED_OBJ VARCHAR)", "is_dttm": False},
{"name": "row_column.nested_obj", "type": "VARCHAR", "is_dttm": False},
{"name": "array_column", "type": "ARRAY(BIGINT)", "is_dttm": False},
{
"column_name": "row_column",
"name": "row_column",
"type": "ROW(NESTED_OBJ VARCHAR)",
"is_dttm": False,
},
{
"column_name": "row_column.nested_obj",
"name": "row_column.nested_obj",
"type": "VARCHAR",
"is_dttm": False,
},
{
"column_name": "array_column",
"name": "array_column",
"type": "ARRAY(BIGINT)",
"is_dttm": False,
},
]
expected_data = [
@@ -225,7 +250,12 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
]
expected_expanded_cols = [
{"name": "row_column.nested_obj", "type": "VARCHAR", "is_dttm": False}
{
"name": "row_column.nested_obj",
"column_name": "row_column.nested_obj",
"type": "VARCHAR",
"is_dttm": False,
}
]
self.assertEqual(actual_cols, expected_cols)
self.assertEqual(actual_data, expected_data)
@@ -240,6 +270,7 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
cols = [
{
"name": "row_column",
"column_name": "row_column",
"type": "ROW(NESTED_OBJ1 VARCHAR, NESTED_ROW ROW(NESTED_OBJ2 VARCHAR))",
"is_dttm": False,
}
@@ -251,17 +282,25 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
expected_cols = [
{
"name": "row_column",
"column_name": "row_column",
"type": "ROW(NESTED_OBJ1 VARCHAR, NESTED_ROW ROW(NESTED_OBJ2 VARCHAR))",
"is_dttm": False,
},
{"name": "row_column.nested_obj1", "type": "VARCHAR", "is_dttm": False},
{
"name": "row_column.nested_obj1",
"column_name": "row_column.nested_obj1",
"type": "VARCHAR",
"is_dttm": False,
},
{
"name": "row_column.nested_row",
"column_name": "row_column.nested_row",
"type": "ROW(NESTED_OBJ2 VARCHAR)",
"is_dttm": False,
},
{
"name": "row_column.nested_row.nested_obj2",
"column_name": "row_column.nested_row.nested_obj2",
"type": "VARCHAR",
"is_dttm": False,
},
@@ -282,14 +321,21 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
]
expected_expanded_cols = [
{"name": "row_column.nested_obj1", "type": "VARCHAR", "is_dttm": False},
{
"name": "row_column.nested_obj1",
"column_name": "row_column.nested_obj1",
"type": "VARCHAR",
"is_dttm": False,
},
{
"name": "row_column.nested_row",
"column_name": "row_column.nested_row",
"type": "ROW(NESTED_OBJ2 VARCHAR)",
"is_dttm": False,
},
{
"name": "row_column.nested_row.nested_obj2",
"column_name": "row_column.nested_row.nested_obj2",
"type": "VARCHAR",
"is_dttm": False,
},
@@ -307,6 +353,7 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
cols = [
{
"name": "row_column",
"column_name": "row_column",
"type": "ROW(NESTED_ROW ROW(NESTED_OBJ VARCHAR))",
"is_dttm": False,
}
@@ -323,16 +370,19 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
expected_cols = [
{
"name": "row_column",
"column_name": "row_column",
"type": "ROW(NESTED_ROW ROW(NESTED_OBJ VARCHAR))",
"is_dttm": False,
},
{
"name": "row_column.nested_row",
"column_name": "row_column.nested_row",
"type": "ROW(NESTED_OBJ VARCHAR)",
"is_dttm": False,
},
{
"name": "row_column.nested_row.nested_obj",
"column_name": "row_column.nested_row.nested_obj",
"type": "VARCHAR",
"is_dttm": False,
},
@@ -363,11 +413,13 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
expected_expanded_cols = [
{
"name": "row_column.nested_row",
"column_name": "row_column.nested_row",
"type": "ROW(NESTED_OBJ VARCHAR)",
"is_dttm": False,
},
{
"name": "row_column.nested_row.nested_obj",
"column_name": "row_column.nested_row.nested_obj",
"type": "VARCHAR",
"is_dttm": False,
},
@@ -383,9 +435,15 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
)
def test_presto_expand_data_with_complex_array_columns(self):
cols = [
{"name": "int_column", "type": "BIGINT", "is_dttm": False},
{
"name": "int_column",
"column_name": "int_column",
"type": "BIGINT",
"is_dttm": False,
},
{
"name": "array_column",
"column_name": "array_column",
"type": "ARRAY(ROW(NESTED_ARRAY ARRAY(ROW(NESTED_OBJ VARCHAR))))",
"is_dttm": False,
},
@@ -398,19 +456,27 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
cols, data
)
expected_cols = [
{"name": "int_column", "type": "BIGINT", "is_dttm": False},
{
"name": "int_column",
"column_name": "int_column",
"type": "BIGINT",
"is_dttm": False,
},
{
"name": "array_column",
"column_name": "array_column",
"type": "ARRAY(ROW(NESTED_ARRAY ARRAY(ROW(NESTED_OBJ VARCHAR))))",
"is_dttm": False,
},
{
"name": "array_column.nested_array",
"column_name": "array_column.nested_array",
"type": "ARRAY(ROW(NESTED_OBJ VARCHAR))",
"is_dttm": False,
},
{
"name": "array_column.nested_array.nested_obj",
"column_name": "array_column.nested_array.nested_obj",
"type": "VARCHAR",
"is_dttm": False,
},
@@ -468,11 +534,13 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
expected_expanded_cols = [
{
"name": "array_column.nested_array",
"column_name": "array_column.nested_array",
"type": "ARRAY(ROW(NESTED_OBJ VARCHAR))",
"is_dttm": False,
},
{
"name": "array_column.nested_array.nested_obj",
"column_name": "array_column.nested_array.nested_obj",
"type": "VARCHAR",
"is_dttm": False,
},
@@ -575,9 +643,20 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
)
def test_presto_expand_data_array(self):
cols = [
{"name": "event_id", "type": "VARCHAR", "is_dttm": False},
{"name": "timestamp", "type": "BIGINT", "is_dttm": False},
{
"column_name": "event_id",
"name": "event_id",
"type": "VARCHAR",
"is_dttm": False,
},
{
"column_name": "timestamp",
"name": "timestamp",
"type": "BIGINT",
"is_dttm": False,
},
{
"column_name": "user",
"name": "user",
"type": "ROW(ID BIGINT, FIRST_NAME VARCHAR, LAST_NAME VARCHAR)",
"is_dttm": False,
@@ -594,16 +673,42 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
cols, data
)
expected_cols = [
{"name": "event_id", "type": "VARCHAR", "is_dttm": False},
{"name": "timestamp", "type": "BIGINT", "is_dttm": False},
{
"column_name": "event_id",
"name": "event_id",
"type": "VARCHAR",
"is_dttm": False,
},
{
"column_name": "timestamp",
"name": "timestamp",
"type": "BIGINT",
"is_dttm": False,
},
{
"column_name": "user",
"name": "user",
"type": "ROW(ID BIGINT, FIRST_NAME VARCHAR, LAST_NAME VARCHAR)",
"is_dttm": False,
},
{"name": "user.id", "type": "BIGINT", "is_dttm": False},
{"name": "user.first_name", "type": "VARCHAR", "is_dttm": False},
{"name": "user.last_name", "type": "VARCHAR", "is_dttm": False},
{
"column_name": "user.id",
"name": "user.id",
"type": "BIGINT",
"is_dttm": False,
},
{
"column_name": "user.first_name",
"name": "user.first_name",
"type": "VARCHAR",
"is_dttm": False,
},
{
"column_name": "user.last_name",
"name": "user.last_name",
"type": "VARCHAR",
"is_dttm": False,
},
]
expected_data = [
{
@@ -616,9 +721,24 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
}
]
expected_expanded_cols = [
{"name": "user.id", "type": "BIGINT", "is_dttm": False},
{"name": "user.first_name", "type": "VARCHAR", "is_dttm": False},
{"name": "user.last_name", "type": "VARCHAR", "is_dttm": False},
{
"column_name": "user.id",
"name": "user.id",
"type": "BIGINT",
"is_dttm": False,
},
{
"column_name": "user.first_name",
"name": "user.first_name",
"type": "VARCHAR",
"is_dttm": False,
},
{
"column_name": "user.last_name",
"name": "user.last_name",
"type": "VARCHAR",
"is_dttm": False,
},
]
self.assertEqual(actual_cols, expected_cols)
@@ -736,12 +856,12 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
table_name = "table_name"
engine = mock.Mock()
cols = [
{"name": "val1"},
{"name": "val2<?!@#$312,/'][p098"},
{"name": ".val2"},
{"name": "val2."},
{"name": "val.2"},
{"name": ".val2."},
{"column_name": "val1"},
{"column_name": "val2<?!@#$312,/'][p098"},
{"column_name": ".val2"},
{"column_name": "val2."},
{"column_name": "val.2"},
{"column_name": ".val2."},
]
PrestoEngineSpec.select_star(
database, table_name, engine, show_cols=True, cols=cols
@@ -756,8 +876,8 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
True,
True,
[
{"name": "val1"},
{"name": "val2<?!@#$312,/'][p098"},
{"column_name": "val1"},
{"column_name": "val2<?!@#$312,/'][p098"},
],
)