SNOW-1222806: MULTI_STATEMENT_COUNT parameter doesn't work in Snowflake SQLAlchemy

[alexgrand]

Please answer these questions before submitting your issue. Thanks!

1. What version of Python are you using?

   Python 3.10.12 (main, Jul 21 2023, 13:21:20) [GCC 11.3.0]

2. What operating system and processor architecture are you using?

   Linux-5.19.0-50-generic-x86_64-with-glibc2.35

3. What are the component versions in the environment (pip freeze)?

   Snowflake connector, sqlalchemy:

   pip freeze | grep snowflake
   snowflake-connector-python==3.0.4
   snowflake-sqlalchemy==1.5.1

4. What did you do?

from sqlalchemy import create_engine
from sqlalchemy.engine import URL
from sqlalchemy.sql import text

engine = create_engine(
    url=URL.create(
        drivername='snowflake',
        username=<username>,
        password=<password>,
        host=<account>,
        query=dict(role="DBA")
    ),
    echo=True,
    echo_pool="debug",
    pool_pre_ping=True,
)

query = text("SET ( var1, var2) = ('test1', 'test2'); SELECT $var1, $var2;")

with engine.connect() as conn:
    conn.execute(query, num_statements=0)

>>>
│ ╭───────────────────────────────── locals ─────────────────────────────────╮ │
│ │      connection = <snowflake.connector.connection.SnowflakeConnection    │ │
│ │                   object at 0x7f7a0a133700>                              │ │
│ │          cursor = <snowflake.connector.cursor.SnowflakeCursor object at  │ │
│ │                   0x7f7a0cd3b0a0>                                        │ │
│ │ done_format_msg = False                                                  │ │
│ │           errno = 8                                                      │ │
│ │     error_class = <class 'snowflake.connector.errors.ProgrammingError'>  │ │
│ │     error_value = {                                                      │ │
│ │                   │   'msg': 'Actual statement count 2 did not match the │ │
│ │                   desired statement count 1.',                           │ │
│ │                   │   'errno': 8,                                        │ │
│ │                   │   'sqlstate': '0A000',                               │ │
│ │                   │   'sfqid': '01b2d431-0002-2d12-0055-4687000210aa',   │ │
│ │                   │   'query': "SET ( var1, var2) = ('test1', 'test2');  │ │
│ │                   SELECT $var1, $var2;",                                 │ │
│ │                   │   'done_format_msg': False                           │ │
│ │                   }                                                      │ │
│ ╰──────────────────────────────────────────────────────────────────────────╯ │
╰──────────────────────────────────────────────────────────────────────────────╯
ProgrammingError: (snowflake.connector.errors.ProgrammingError) 000008 (0A000): 
Actual statement count 2 did not match the desired statement count 1.
[SQL: SET ( var1, var2) = ('test1', 'test2'); SELECT $var1, $var2;]

5. What did you expect to see?

   I expect to be able to run several statements in one SQL statement from file. Snowflake SQLAlchemy works same as python snowflake connector ^version 2.9.0

[alexgrand]

Looks like SQLAlchemy runs sqlalchemy.engine.default.DefaultDialect.do_execute method which doesn't pass num_statements parameter to the snowflake.connector.cursor.SnowflakeCursor.execute and num_statements is always set to None

[alexgrand]

Fixed by rewriting sqlalchemy cursor command via sqlalchemy do_execute event:

from sqlalchemy import event

@event.listens_for(<engine>, 'do_execute')
def do_execute(cursor, statement, parameters, context):
    return cursor.execute(statement, parameters, num_statements=0)

Any better option to fix it? Maybe I miss something?

[sfc-gh-sghosh]

Hello @alexgrand ,

Thank you for bringing up this matter. Currently, the MULTI_STATEMENT_COUNT parameter isn't supported for Snowflake SQLAlchemy. However, there are plans to tentatively support it in Q2 2024.

Regards, Sujan

[mrcolumbia]

Has there been any movement on adding this feature? In my understanding, Snowflake performs query planning optimization when presented with multiple SQL statements at once - this optimization would be lost if queries are submitted one at a time. I believe that makes this feature essential.

[mrcolumbia]

Fixed by rewriting sqlalchemy cursor command via sqlalchemy do_execute event:

from sqlalchemy import event

@event.listens_for(<engine>, 'do_execute')
def do_execute(cursor, statement, parameters, context):
    return cursor.execute(statement, parameters, num_statements=0)

Any better option to fix it? Maybe I miss something?

This works for me also, but did you find any way of doing performance logging on each of the individual queries in the multi-query file?