Using CASE in a select statement breaks the split_statements logic used for inferring lineage from stored procedures in mssql

[ttekampe]

Describe the bug
Using CASE in a select statement breaks the split_statements logic used for inferring lineage from stored procedures in mssql

To Reproduce

from datahub.sql_parsing.schema_resolver import SchemaResolver
from datahub.sql_parsing.split_statements import split_statements
from datahub.sql_parsing.sqlglot_lineage import sqlglot_lineage

sql_text = """
SELECT
    Id,
    CASE
        WHEN Id > 10 THEN 'IS_GREATER_THAN_TEN'
        WHEN Id > 5 THEN 'IS_GREATER_THAN_FIVE'
        ELSE 'IS_SMALL'
    END AS foo,
    bar
FROM mySchema.myTable
"""

schema_resolver = SchemaResolver(platform="mssql")

statements = split_statements(sql_text)
results = []
for s in statements:
    print(s)
    print("------")
    parsed = sqlglot_lineage(
        s,
        schema_resolver=schema_resolver,
        default_db="myDatabase",
        default_schema="mySchema",
    )
    results.append(parsed)


for r in results:
    print(r.json(indent=4))
    print("------")

The individual statements that are produced by split_statements are

SELECT
    Id,
    CASE
        WHEN Id > 10 THEN 'IS_GREATER_THAN_TEN'
        WHEN Id > 5 THEN 'IS_GREATER_THAN_FIVE'
        ELSE 'IS_SMALL'
------
END
------
AS foo,
    bar
FROM mySchema.myTable
------

Expected behavior
The statement should not be broken up, such that the SELECT part can be parsed correctly and lineage can be inferred.

Additional context
acryl-datahub python package version is 0.15.0.5

[ttekampe]

I wanted to prepare a PR just to find out that this is already fixed on the master branch.

Any plans to release a new version of the acryl-datahub python package, so we can benefit from these improvements?

[ttekampe]

There is still one thing that leaves me a bit puzzled:

The regex patterns to match control flow keywords here

datahub/metadata-ingestion/src/datahub/sql_parsing/split_statements.py

Lines 9 to 23 in 7472c53

	CONTROL_FLOW_KEYWORDS = [
	"GO",
	r"BEGIN\w+TRY",
	r"BEGIN\w+CATCH",
	"BEGIN",
	r"END\w+TRY",
	r"END\w+CATCH",
	# This isn't strictly correct, but we assume that IF | (condition) | (block) should all be split up
	# This mainly ensures that IF statements don't get tacked onto the previous statement incorrectly
	"IF",
	# For things like CASE, END does not mean the end of a statement.
	# We have special handling for this.
	END_KEYWORD,
	# "ELSE", # else is also valid in CASE, so we we can't use it here.
	]

Why do you use \w+ and not \s+?

For example BEGIN\w+TRY matches BEGINFooBarTRY, but it does not match BEGIN TRY

@hsheth2 and @mayurinehate is this intentional or simply a typo?

[hsheth2]

Yeah that looks like a typo - we'd welcome a PR with some corresponding tests :)