test(ingest): test case statements with sql parser (#8437)

This commit is contained in:
Harshal Sheth 2023-08-01 07:04:48 -07:00 committed by GitHub
parent ef3b9489aa
commit 66074341f7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 65 additions and 0 deletions

View File

@ -0,0 +1,37 @@
{
"query_type": "SELECT",
"in_tables": [
"urn:li:dataset:(urn:li:dataPlatform:snowflake,snowflake_sample_data.tpch_sf1.orders,PROD)"
],
"out_tables": [],
"column_lineage": [
{
"downstream": {
"table": null,
"column": "TOTAL_PRICE_CATEGORY"
},
"upstreams": [
{
"table": "urn:li:dataset:(urn:li:dataPlatform:snowflake,snowflake_sample_data.tpch_sf1.orders,PROD)",
"column": "totalprice"
}
]
},
{
"downstream": {
"table": null,
"column": "TOTAL_PRICE_SUCCESS"
},
"upstreams": [
{
"table": "urn:li:dataset:(urn:li:dataPlatform:snowflake,snowflake_sample_data.tpch_sf1.orders,PROD)",
"column": "is_payment_successful"
},
{
"table": "urn:li:dataset:(urn:li:dataPlatform:snowflake,snowflake_sample_data.tpch_sf1.orders,PROD)",
"column": "totalprice"
}
]
}
]
}

View File

@ -287,6 +287,34 @@ FROM snowflake_sample_data.tpch_sf1.orders o
)
def test_snowflake_case_statement():
assert_sql_result(
"""
SELECT
CASE
WHEN o."totalprice" > 1000 THEN 'high'
WHEN o."totalprice" > 100 THEN 'medium'
ELSE 'low'
END as total_price_category,
-- Also add a case where the column is in the THEN clause.
CASE
WHEN o."is_payment_successful" THEN o."totalprice"
ELSE 0
END as total_price_success
FROM snowflake_sample_data.tpch_sf1.orders o
""",
dialect="snowflake",
schemas={
"urn:li:dataset:(urn:li:dataPlatform:snowflake,snowflake_sample_data.tpch_sf1.orders,PROD)": {
"orderkey": "NUMBER",
"totalprice": "FLOAT",
"is_payment_successful": "BOOLEAN",
},
},
expected_file=RESOURCE_DIR / "test_snowflake_case_statement.json",
)
@pytest.mark.skip(reason="We don't handle the unnest lineage correctly")
def test_bigquery_unnest_columns():
assert_sql_result(