python中的上下文管理器如何捕获异常

Question

I have the following method to setup a connection with the DB and tear it down in the end. Function looks something like this

def read_db(self, sql_statement):

    conn = pymysql.connect(host=self.h,user=self.u,passwd=self.pw,
                           db=self.db,port=self.p)
    try:
        with conn.cursor() as cur:
            cur.execute(sql_statement)
            doStuffGeneratingException()

    except Exception:
        cur.rollback()
        raise

    finally:
        conn.close()   

Now if I had to replace this with context manager, I think it would look something like

@contextmanager
def setup_sql():
    conn = pymysql.connect(host=self.h,user=self.u,passwd=self.pw,
                           db=self.db,port=self.p)
    yield conn.cursor()
    connection.cursor().close()
    self.connection.close()


def read_db(self, sql_statement):
    with setup_sql() as cur:
    try:
        cur.execute(sql_statement)
        doStuffGeneratingException()

    except:
         cur.rollback()
         raise

Now my questions are

1. Is the above interpretation of context manager correct?
2. If there was an error that occurred while doing pymysql.connect(...) statement inside contextmanager , how would that be handled? How would it be bubbled up to the calling function?
3. What happens if there is an exception in doStuffGeneratingException() in the with implementation? Will the control go first to the setup_sql to execute the statements after yield ?

Answer 1

1, sorta. The entire try/except needs another level of indentation.

def read_db(self, sql_statement):
    with setup_sql() as cur:
        try:
            cur.execute(sql_statement)
            doStuffGeneratingException()

        except:
             cur.rollback()
             raise

2, An error there isn't handled anywhere in your code, so python itself would report on the exception and halt execution. It could be caught anywhere you chose. Inside setup_sql() and read_db() are both viable, but typically you want to handle exceptions as close to what's raising them as possible if you intend on doing something about it. To do it inside read_db() another try: block would be needed around your with setup_sql():

def read_db(self, sql_statement):
    try:
        with setup_sql() as cur:
            try:
                cur.execute(sql_statement)
                doStuffGeneratingException()

            except:
                 # gets exceptions thrown by cur.execute() and doStuffGeneratingException() 
                 # will not process context manager statements after yield if flow doesn't continue in this function past this except block
                 cur.rollback()
                 raise
    except:
        # gets exceptions thrown by `with setup_sql() as cur:`
        # again none of the statements within the context manager after the statement that raises an exception will be executed
        pass

3, no. An exception is an immediate 'return' it will hit your rollback, and re-raising it will abort your with block. If you want the context manager to complete, catch the exceptions and deal with them without re-raising. If you need to raise an exception there and want the context manager to finish its work, set a flag variable and raise after closure of the context manager or refactor your code another way to achieve that objective.

Answer 2

I believe you actually can encapsulate the error handling in the context manager by placing try / except / finally around the yield statement inside of the method, like:

from contextlib import contextmanager

@contextmanager
def setup_sql():
    conn = pymysql.connect(host=self.h,user=self.u,passwd=self.pw,
                           db=self.db,port=self.p)
    cursor = conn.cursor()
    try:
        yield cursor
    except Exception:
        cursor.rollback()
        raise
    finally:
        cursor.close()
        conn.close()

def read_db(self, sql_statement):
    with setup_sql() as cur:
        cur.execute(sql_statement)
        doStuffGeneratingException()

I haven't tried this, but I had come across this comment on another SO question that links to the documentation about @contextmanager explaining how this works.