SNOW-782479: Context manager for Transactions

[jamesweakley]

What is the current behavior?

Currently the user is required to keep track of what state transactions are in, and commit/rollback as needed based on different execution paths. This is particularly important when stored procs are nested, because you don't want an open transaction crossing the proc boundary.

Code inside a stored proc might look like this:

transaction_open = False
try:
  session.sql('begin transaction').collect()
  transaction_open = True
  session.sql('truncate table my_table').collect()
  # some other code that might error
  # ...
  session.sql('insert into my_table(col1) values(1)').collect()
  session.sql('commit').collect()
  transaction_open = False
  # some other stuff after the transactions (maybe more transactions?)
  # this might also error
  return {
    "success" : True
  }
except Exception as exception:
  if transaction_open is True:
    session.sql('rollback').collect()
  return {
    "success": False,
    "error": str(exception)
  }

What is the desired behavior?

Something more pythonic via a context manager, similar to the way you can manage connections.

The above code could look more like this:

try:
  with session.transaction(action_on_complete='commit',action_on_error='rollback') as txn:
    session.sql('truncate table my_table').collect()
    # some other code that might error
    # ...
    session.sql('insert into my_table(col1) values(1)').collect()
    # some other stuff after the transactions (maybe more transactions?)
    # this might also error
    return {
      "success" : True
    }
except Exception as exception:
  return {
    "success": False,
    "error": str(exception)
  }

Internally it would be running BEGIN TRANSACTION etc for you in the __enter__ and __exit__ methods. You could also override behaviour, e.g. manual rollback with txn.rollback() followed by a return statement and the context manager then wouldn't do the commit when it exits.

How would this improve snowflake-snowpark-python?

It would allow people to use transactions with less cognitive overhead, less room for bugs, and more terse code.

References, Other Background

I am happy to contribute a PR, but it would be good to first know if people agree with the concept or have any feedback/concerns.

[sfc-gh-achandrasekaran]

Thanks @jamesweakley . We will discuss internally and get back to you on this

[gbatiz]

Would be great to have this, less cognitive load, less error-prone, follows the de-facto standard SQLAlchemy-logic.

[jamesweakley]

FWIW, I'm currently using this class as a context manager:

from snowflake.snowpark import Session
from typing import Literal
SnowflakeTransactionAction = Literal["commit", "rollback"]


class SnowflakeTransaction:
    def __init__(
        self,
        session: Session,
        action_on_complete: SnowflakeTransactionAction = "commit",
        action_on_error: SnowflakeTransactionAction = "rollback",
    ):
        self.session = session
        self.actioned = False
        self.action_on_complete = action_on_complete
        if action_on_complete not in ["commit", "rollback"]:
            raise ValueError(f"Invalid action_on_complete action {action_on_complete}")
        self.action_on_error = action_on_error
        if action_on_error not in ["commit", "rollback"]:
            raise ValueError(f"Invalid action_on_error action {action_on_error}")

    def __enter__(self):
        self.session.sql("begin transaction").collect()
        return self

    def commit(self):
        self.session.sql("commit").collect()
        self.actioned = True

    def rollback(self):
        self.session.sql("rollback").collect()
        self.actioned = True

    def __exit__(self, exc_type, exc_val, exc_tb):
        # if we already actioned, don't do anything
        if not self.actioned:
            # if an error was thrown, rollback
            if exc_type is not None:
                self.session.sql(self.action_on_error).collect()
            else:
                self.session.sql(self.action_on_complete).collect()

Usage:

with SnowflakeTransaction(session):
  # sql queries, etc

[gbatiz]

Thanks for sharing this!

[shabazy]

__exit__

I like this solution. I've tested and debugged this carefully and working fine. Thanks.