SnowflakeHook.run() doesn't accept multiple statements without split_statements

[set92]

Apache Airflow version

Other Airflow 2 version (please specify below)

If "Other Airflow 2 version" selected, which one?

2.10.3

What happened?

I have a query like this:

BEGIN;
CREATE TABLE ();
INSERT INTO () SELECT FROM ();
INSERT INTO () SELECT FROM ();
DELETE FROM WHERE;
COMMIT;

If using SnowflakeSqlApiOperator() runs without any problem, or without specifying anything. All in a single block. With the hook I tried to run it with semicolons, without them (snowflake support team recommended it), and the only option that worked was with the split_statements=True, but then you have multiple query id, harder to review, and I think the transaction will keep intact, but not sure.

Also if you try to run with split_statements=False, it will throw the error

Traceback (most recent call last):
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/airflow/models/taskinstance.py", line 767, in _execute_task
    result = _execute_callable(context=context, **execute_callable_kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/airflow/models/taskinstance.py", line 733, in _execute_callable
    return ExecutionCallableRunner(
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/airflow/utils/operator_helpers.py", line 252, in run
    return self.func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/airflow/models/baseoperator.py", line 417, in wrapper
    return func(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/airflow/operators/python.py", line 238, in execute
    return_value = self.execute_callable()
                   ^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/airflow/operators/python.py", line 256, in execute_callable
    return runner.run(*self.op_args, **self.op_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/airflow/utils/operator_helpers.py", line 252, in run
    return self.func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/airflow/dags/ghs/bi/spl_on_snowflake/actions.py", line 71, in create_pu_per_client_func
    snflk_hook.run(sql_statement, split_statements=False)
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/airflow/providers/snowflake/hooks/snowflake.py", line 435, in run
    self._run_command(cur, sql_statement, parameters)  # type: ignore[attr-defined]
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/airflow/providers/common/sql/hooks/sql.py", line 549, in _run_command
    cur.execute(sql_statement)
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/snowflake/connector/cursor.py", line 1097, in execute
    Error.errorhandler_wrapper(self.connection, self, error_class, errvalue)
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/snowflake/connector/errors.py", line 284, in errorhandler_wrapper
    handed_over = Error.hand_to_other_handler(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/snowflake/connector/errors.py", line 339, in hand_to_other_handler
    cursor.errorhandler(connection, cursor, error_class, error_value)
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/snowflake/connector/errors.py", line 215, in default_errorhandler
    raise error_class(
snowflake.connector.errors.ProgrammingError: 000008 (0A000): 01bb3de5-020c-cab0-0006-21fb059c3996: Actual statement count 6 did not match the desired statement count 1.

Following with this error you can get to https://docs.snowflake.com/en/developer-guide/sql-api/submitting-multiple-statements, where it mentions a parameter MULTI_STATEMENT_COUNT that can be modified to send the same number as the queries you have in the string and it will work. But I wasn't sure how, or if it is possible to pass it to the Hook.

What you think should happen instead?

As with SnowflakeSqlApiOperator(), I would have expected to be able to run all in a single string.

How to reproduce

I suppose you can test it with

sql_statement = """
BEGIN TRANSACTION;
CREATE TABLE Persons (id int);
SELECT * FROM Persons;
COMMIT;
"""
snflk_hook = SnowflakeHook(conn_id=snflk_conn_id).run(sql_statement, split_statements=False)

Operating System

MWAA - Linux

Versions of Apache Airflow Providers

No response

Deployment

Amazon (AWS) MWAA

Deployment details

No response

Anything else?

No response

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[set92]

I have unlocked my problem using the SnowflakeSqlApiHook(conn_id=snflk_conn_id).execute_query(sql) what is what SnowflakeSqlApiOperator() uses, but this one needs a JWT Token, so I'm not liking that much.

And also, then the split_statements parameter in SnowflakeHook() doesn't make sense. Because if you set it to false, it only works for one statement, and therefore doesn't matter split_statements=True or False. So, would be nice to have MULTI_STATEMENT_COUNT parameter to be able to use it with split_statement=False.

EDIT:

I was able to use the MULTI_STATEMENT_COUNT like snflk_hook = SnowflakeHook(conn_id=snflk_conn_id, session_parameters={"MULTI_STATEMENT_COUNT": 6}), but the error I get is the next one:

[2025-03-25, 11:09:30 UTC] {taskinstance.py:3311} ERROR - Task failed with exception
Traceback (most recent call last):
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/airflow/models/taskinstance.py", line 767, in _execute_task
    result = _execute_callable(context=context, **execute_callable_kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/airflow/models/taskinstance.py", line 733, in _execute_callable
    return ExecutionCallableRunner(
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/airflow/utils/operator_helpers.py", line 252, in run
    return self.func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/airflow/models/baseoperator.py", line 417, in wrapper
    return func(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/airflow/operators/python.py", line 238, in execute
    return_value = self.execute_callable()
                   ^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/airflow/operators/python.py", line 256, in execute_callable
    return runner.run(*self.op_args, **self.op_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/airflow/utils/operator_helpers.py", line 252, in run
    return self.func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/airflow/dags/ghs/bi/spl_on_snowflake/actions.py", line 71, in create_pu_per_client_func
    snflk_hook.run(sql_statement, split_statements=False)
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/airflow/providers/snowflake/hooks/snowflake.py", line 435, in run
    self._run_command(cur, sql_statement, parameters)  # type: ignore[attr-defined]
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/airflow/providers/common/sql/hooks/sql.py", line 549, in _run_command
    cur.execute(sql_statement)
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/snowflake/connector/cursor.py", line 1097, in execute
    Error.errorhandler_wrapper(self.connection, self, error_class, errvalue)
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/snowflake/connector/errors.py", line 284, in errorhandler_wrapper
    handed_over = Error.hand_to_other_handler(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/snowflake/connector/errors.py", line 339, in hand_to_other_handler
    cursor.errorhandler(connection, cursor, error_class, error_value)
  File "/usr/local/airflow/.local/lib/python3.11/site-packages/snowflake/connector/errors.py", line 215, in default_errorhandler
    raise error_class(
snowflake.connector.errors.ProgrammingError: 100132 (P0000): 01bb3e1d-020c-b678-0006-21fb059c6aca: JavaScript execution error: Uncaught Execution of multiple statements failed on statement "BEGIN TRANSACTION" (at line 1, position 0).
Actual statement count 1 did not match the desired statement count 6. in SYSTEM$MULTISTMT at '    throw `Execution of multiple statements failed on statement {0} (at line {1}, position {2}).`.replace('{1}', LINES[i])' position 4
stackstrace: 
SYSTEM$MULTISTMT line: 10

[kunaljubce]

I am triaging this, will update hopefully with a fix!

[kunaljubce]

Ok @set92, I debugged this and it seems like the difference in behaviour between SnowflakeHook and SnowflakeSqlApiOperator is exactly as intended with regards to how they treat multi-sql blocks. And this is partly attributed to the way they are authenticating and communicating with Snowflake, i.e. SnowflakeHook allows all forms of authentication including username/password etc, whereas SnowflakeSqlApiOperator only supports Snowflake SQL API authentication via key/pair.

Looking into the Snowflake doc which you have referred above, it also mentions communicating with Snowflake via API for multi-sql statements to be processed as a single query block. This is where SnowflakeHook gets its restriction and for your usecase, it seems prudent to choose either of the operators depending on what you need the most - Single transaction debuggability (SnowflakeSqlApiOperator) or dag code simplicity (SnowflakeHook).

@potiuk Please feel free to add if I am missing something or correct me if I am wrong. Thanks!

[set92]

@kunaljubce Sorry that I didn't have the email notifications ON, and I kept going since I didn't think this issue was going to get to a fix. While working with Snowflake + Airflow I discovered basically what you are saying.

In my mind, when you run a transaction is because you want atomicity, right? Well, Snowflake doesn't think like that. It is needed a change on the parameters, and then use the SnowflakeSqlApiOperator to make sure it will run all sentences or none. For example, the next SQL runs, but needs TRANSACTION_ABORT_ON_ERROR because if not it will commit whatever had been run before crashing.

    sql_statement = """
    ALTER SESSION SET TRANSACTION_ABORT_ON_ERROR = TRUE;
    BEGIN TRANSACTION;

    CREATE OR ALTER TABLE db_test.schema_test.test1 (
        txt character varying(1)
    );

    SELECT * FROM db. ;

    COMMIT;
    """
    snflk_hook = SnowflakeSqlApiHook(snowflake_conn_id=snflk_conn_id)
    snflk_hook.execute_query(sql_statement, statement_count=0)

And of course, use the SnowflakeSqlApiHook().execute_query() instead of SnowflakeHook().run() because if not it will not respect the transactions. I thought it was a nice thing to comment and know about. I feel going to SQLExecuteQueryOperator was nice to standardize all the different databases, but when you see these type of problems, you remember that each database has their nuisances and why is good to respect their way of running or doing things on their own.

[kunaljubce]

@set92 Wow that's really interesting about the atomicity thing.

Just to clarify regarding your statement below, you mean SnowflakeSqlApiHook().execute_query() does NOT respect transactions by default, am I right?

And of course, use the SnowflakeSqlApiHook().execute_query() instead of SnowflakeHook().run() because if not it will not respect the transactions.

In case you tried it out, what about SnowflakeHook().run()? Does this also have the same limitations with transactions and atomicity?

[set92]

@kunaljubce The only one which respect transactions is SnowflakeSqlApiHook(). SnowflakeHook() I feel runs each query in its own session, so it doesn't respect the transactionality.

[shahar1]

Tagging this issue as as both:

• "documentation": to document the behavior by current design.
• "feature request": to enhance the intended behaivor of Snowflake when working with SQLExecuteQueryOperator .

Anyone who wants to do any of the above, feel free to chime in.