SNOW-1662210: Add ability to patch for Session sql method

[thannam-conga]

What is the current behavior?

Unable to patch Session.sql() method. After following the documentation receiving the error: NotImplementedError: [Local Testing] Session.sql is not supported.'

import pytest
from project.utils import get_env_var_config
from functools import partial
from snowflake.snowpark.session import Session

def pytest_addoption(parser):
    parser.addoption("--snowflake-session", action="store", default="live")

@pytest.fixture(scope='module')
def session(request) -> Session:
    if request.config.getoption('--snowflake-session') == 'local':
        return Session.builder.configs({'local_testing': True}).create()
    else:
        return Session.builder.configs(get_env_var_config()).create()

Error: NotImplementedError: [Local Testing] Session.sql is not supported.

What is the desired behavior?

Ability to patch Session.sql() so code that uses Session.sql() can have tests made

# test\conftest.py
import pytest
from project.utils import get_env_var_config
from snowflake.snowpark.session import Session
from unittest.mock import patch

def pytest_addoption(parser):
    parser.addoption("--snowflake-session", action="store", default="live")

def mock_sql():
    def mock_sql(session, sql_string):  # patch for SQL operations
        if sql_string == "select 1,2,3":
            return session.create_dataframe([[1,2,3]])
        else:
            raise RuntimeError(f"Unexpected query execution: {sql_string}")

@pytest.fixture(scope='module')
def session(request) -> Session:
    if request.config.getoption('--snowflake-session') == 'local':
        mock_session = Session.builder.configs({'local_testing': True}).create()

        with patch.object(mock_session, 'sql', wraps=partial(mock_sql, mock_session):
            return mock_session

    else:
        return Session.builder.configs(get_env_var_config()).create()

# test\test_sproc.py
def test_sproc(session):
    DB = 'CITIBIKE'
    SCHEMA = 'TEST'

    # Set up source table
    tbl = session.create_dataframe(
        data=[
            [1983, '2018-03-01 09:47:00.000 +0000', 551, 30958],
            [1988, '2018-03-01 09:47:01.000 +0000', 242, 19278],
            [1992, '2018-03-01 09:47:01.000 +0000', 768, 18461],
            [1980, '2018-03-01 09:47:03.000 +0000', 690, 15533],
            [1991, '2018-03-01 09:47:03.000 +0000', 490, 32449],
            [1959, '2018-03-01 09:47:04.000 +0000', 457, 29411],
            [1971, '2018-03-01 09:47:08.000 +0000', 279, 28015],
            [1964, '2018-03-01 09:47:09.000 +0000', 546, 15148],
            [1983, '2018-03-01 09:47:11.000 +0000', 358, 16967],
            [1985, '2018-03-01 09:47:12.000 +0000', 848, 20644],
            [1984, '2018-03-01 09:47:14.000 +0000', 295, 16365]
        ],
        schema=['BIRTH_YEAR', 'STARTTIME', 'TRIPDURATION',    'BIKEID'],
    )

    tbl.write.mode('overwrite').save_as_table([DB, SCHEMA, 'TRIPS_TEST'], mode='overwrite')

    # Stored procedure that will use sql method multiple times
    sproc(session)
    
    assert final_table == expected_final_table  # will assert that the entire process works by evaluating the final output

If this is not an existing feature in snowflake-snowpark-python. How would this impact/improve non-local testing mode?

It would allow users to work around the NotImplementedError for Session.sql when it is required in the code

References, Other Background

https://docs.snowflake.com/en/developer-guide/snowpark/python/tutorials/testing-tutorial#configure-local-testing

[sfc-gh-aling]

thanks for reaching out.

for sql operation, presently we recommend patching the sql method manually just like what you posted: https://docs.snowflake.com/en/developer-guide/snowpark/python/testing-locally#sql-operations

may I ask if you are looking for other ways to patch the sql call?

[thannam-conga]

I don't think I am doing it correctly. What I am trying to accomplish is patching the sql method to create the desired dataframes for testing a sproc. However, during setup I get the same error NotImplementedError: [Local Testing] Session.sql is not supported. I have updated the code to reflect what I am trying

After reviewing another issue I found this works...

@pytest.fixture(scope='module')
def session(request) -> Session:
    if request.config.getoption('--snowflake-session') == 'local':
        Session.sql = mock_sql
        return = Session.builder.configs({'local_testing': True}).create()

    else:
        return Session.builder.configs(get_env_var_config()).create()

Not sure this will have unintended consequences but I will plan on using this approach