Extract set values from UPDATE statement

[vittxr]

I have some MySQL update statements:

sql_query = """UPDATE Customers
               SET ContactName = 'Alfred Schmidt', City = 'Frankfurt'
               WHERE CustomerID = 1 AND CustomerName = 'test'"""

sql_query = """UPDATE Customers
               SET is_deleted = 0
               WHERE CustomerID = 10"""

I extracted the set values successfully using this function:

def extract_set_values(tokens) -> dict[str, str]:
    set_values = {}
    for token in tokens:
        if isinstance(token, sqlparse.sql.Comparison):
            column = token.left.get_real_name()
            value = token.right.value.strip("'")
            set_values[column] = value

        if isinstance(token, sqlparse.sql.IdentifierList):
            set_values = extract_set_values(token.tokens)

    return set_values

This is not good in my opinion because Comparison refers to WHERE comparison statements. The correct way to do this should be with the Assignment class, but it is just for comparisons with :=. If I change the SQL query to:

sql_query = """UPDATE Customers
               SET ContactName := 'Alfred Schmidt', City := 'Frankfurt'
               WHERE CustomerID = 1 AND CustomerName = 'test'"""

then I could write the function as:

def extract_set_values(tokens) -> dict[str, str]:
    set_values = {}
    for token in tokens:
        if isinstance(token, sqlparse.sql.Assignment):
            splited = token.value.split(":=")
            column, value = splited[0], splited[1]
            set_values[column] = value

    return set_values

This seems more appropriate. Is it difficult to mark ContactName = 'Alfred Schmidt' as Assignment instead of Comparison if it has a SET prefix?