Vim - SQL highlighting in Python strings

python.vim

" Put this in ~/.vim/after/syntax/python.vim

" Slight tweaks to https://lonelycoding.com/can-i-use-both-python-and-sql-syntax-highlighting-in-the-same-file-in-vim/

" Needed to make syntax/sql.vim do something
unlet b:current_syntax

" Load SQL syntax
syntax include @SQL syntax/sql.vim

" Need to add the keepend here
syn region pythonString matchgroup=pythonQuotes
      \ start=+[uU]\=\z(['"]\)+ end="\z1" skip="\\\\\|\\\z1"
      \ contains=pythonEscape,@Spell keepend
syn region  pythonRawString matchgroup=pythonQuotes
      \ start=+[uU]\=[rR]\z(['"]\)+ end="\z1" skip="\\\\\|\\\z1"
      \ contains=@Spell keepend

syn region SQLEmbedded contains=@SQL containedin=pythonString,pythonRawString contained
    \ start=+\v(ALTER|BEGIN|CALL|COMMENT|COMMIT|CONNECT|CREATE|DELETE|DROP|END|EXPLAIN|EXPORT|GRANT|IMPORT|INSERT|LOAD|LOCK|MERGE|REFRESH|RENAME|REPLACE|REVOKE|ROLLBACK|SELECT|SET|TRUNCATE|UNLOAD|UNSET|UPDATE|UPSERT)+
    \ end=+;+

let b:current_syntax = "pysql"

This is great! One thing I modified was adding -- (start=+\v(--|ALTER|BEGIN|...) to catch strings that start with a comment.

The other thing I can't seem to figure out how to solve is it doesn't work for me with arrays containing triple quotes. For example:

# works fine
x = """
    -- this works fine
    select * from x;
"""

# also works
y = ['select * from y;', "-- still works"

# doesn't work
z = [
    """
        select * from z;
    """,
    """
        -- nope, doesn't work
        insert into z values (0, null);
    """
]