This commit is contained in:
GerbenvdHuizen 2024-05-05 02:07:59 -03:00 committed by GitHub
commit 6cb70f2b7e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 35 additions and 0 deletions

View File

@ -218,6 +218,7 @@ class HiveEngineSpec(PrestoEngineSpec):
np.dtype("float64"): "DOUBLE",
np.dtype("int64"): "BIGINT",
np.dtype("object"): "STRING",
np.dtype("datetime64[ns]"): "TIMESTAMP",
}
return hive_type_by_dtype.get(dtype, "STRING")

View File

@ -221,7 +221,41 @@ def test_df_to_sql_if_exists_replace_with_schema(mock_upload_to_s3, mock_g):
mock_execute.assert_any_call(f"DROP TABLE IF EXISTS {schema}.{table_name}")
app.config = config
@mock.patch("superset.db_engine_specs.hive.g", spec={})
@mock.patch("superset.db_engine_specs.hive.upload_to_s3")
def test_df_to_sql_schema_definition(mock_upload_to_s3, mock_g):
config = app.config.copy()
app.config["CSV_TO_HIVE_UPLOAD_DIRECTORY_FUNC"]: lambda *args: ""
mock_location = "mock-location"
mock_upload_to_s3.return_value = mock_location
mock_g.user = True
mock_database = mock.MagicMock()
mock_database.get_df.return_value.empty = False
mock_execute = mock.MagicMock(return_value=True)
mock_database.get_sqla_engine.return_value.execute = mock_execute
table_name = "foobar"
df = pd.DataFrame({
"bool": pd.Series(dtype="bool"),
"float64": pd.Series(dtype="float64"),
"int64": pd.Series(dtype="int64"),
"object": pd.Series(dtype="object"),
"datetime": pd.Series(dtype="datetime64[ns]")
})
expected_schema_definition = "`bool` BOOLEAN, `float64` DOUBLE, `int64` BIGINT, `object` STRING, `datetime` TIMESTAMP"
with app.app_context():
HiveEngineSpec.df_to_sql(
mock_database,
Table(table=table_name),
df,
{"if_exists": "replace"},
)
_, call_args, _ = mock_execute.mock_calls[1]
assert expected_schema_definition in str(call_args[0])
app.config = config
def test_is_readonly():
def is_readonly(sql: str) -> bool:
return HiveEngineSpec.is_readonly_query(ParsedQuery(sql))