fix(staging): recognize streaming input after SQL comments

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>

Author: aakash-saravanan-db

src/databricks/sql/backend/kernel/client.py

@@ -572,9 +572,9 @@ def execute_command(
             try:
                 stripped_sql = _strip_leading_sql_comments(operation)
                 staging_verb = _staging_verb(stripped_sql)
-                if re.match(
-                    r"""PUT\s*(['"])__input_stream__\1""",
-                    stripped_sql,
+                if staging_verb == "PUT" and re.match(
+                    r"""(['"])__input_stream__\1""",
+                    _strip_leading_sql_comments(stripped_sql[3:]),
                     re.IGNORECASE,
                 ):
                     raise NotSupportedError(

tests/unit/test_kernel_client.py

@@ -953,6 +953,8 @@ def test_staging_verb_matches_only_the_leading_statement(operation, expected):
         "/* upload */ PUT '__input_stream__' INTO '/Volumes/c/s/v/file'",
         "/* outer /* inner */ outer */ PUT '__input_stream__' INTO '/Volumes/c/s/v/file'",
         "PUT'__input_stream__' INTO '/Volumes/c/s/v/file'",
+        "PUT /* upload */ '__input_stream__' INTO '/Volumes/c/s/v/file'",
+        "PUT\n-- upload\n'__input_stream__' INTO '/Volumes/c/s/v/file'",
         "/* upload */ PUT\"__input_stream__\" INTO '/Volumes/c/s/v/file'",
     ],
 )