Context managers in dol

[thorwhalen]

Context managers pop up a lot in data interactions. Two frequent cases:

• the data source needs to be "connected" to (and disconnected when finished), such as in most DBs or remote data sources
• we want to do batch operations, such as accumulate a bunch of write operations and send them all at once to the DB

It's our abstraction's job to remove that concern from the user's view as much as possible.
The thing though, is that this may not always be possible without knowing more of the user's intent or context.
If a store sees a single write operation, how does it know if it needs to be executed immediately, or buffered and written later with a bunch of other accumulated operations?

In mongodol, we've developed some general tools in the tracking_methods module (todo: should probably be refactored to include some base functionality in dol or xdol), and a similar "command pattern" appears again in tabled, where accumulation of commands and execution are separated.

[thorwhalen]

Towards an Idempotent Connection Context Manager

The following content can also be found in this notebook. This notebook is also more maintained than the following comment.

Note that a more general reusable tool for managing nested contexts was attempted in the Not entering/exiting contexts twice when nested issue in lkj, but failed for now, so was closed.

Following are more (data bases) specific implementation notes:

Say we have a connection class (or function/method) and a DAO using it:

class Connection:
    def __init__(self, uri='DFLT_URI'):
        self.uri = uri

    def __enter__(self):
        print(f"Opening connection to {self.uri}")
        return self

    def __exit__(self, *exc_info):
        print(f"Closing connection to {self.uri}")

    open = __enter__
    close = __exit__

class DAO1:
    def __init__(self, uri='DFLT_URI'):
        self.connection = Connection(uri)

    def read(self, k):
        with self.connection:
            print(f"    Reading {k}")

dao1 = DAO1()
dao1.read('bob')
dao1.read('alice')
print("")
# Prints:
# Opening connection to DFLT_URI
#     Reading bob
# Closing connection to DFLT_URI
# Opening connection to DFLT_URI
#     Reading alice
# Closing connection to DFLT_URI
#

If we now make our DAO a context manager itself, we get exactly the same behavior (of opening and closing the context at every read). Additionally, we opening the connection twice in a row in the beginning, and close it twice in a row at the end. That could be a problem if our connection object doesn't like being opened when already open, or closed when already closed.

class DAO2(DAO1):
    def __enter__(self):
        print(f"DAO entry...")
        self.connection.__enter__()
        return self

    def __exit__(self, *exc_info):
        print(f"... DAO exit")
        self.connection.__exit__(*exc_info)
        
dao2 = DAO2()
dao2.read('bob')
dao2.read('alice')
print("")
# Opening connection to DFLT_URI
#     Reading bob
# Closing connection to DFLT_URI
# Opening connection to DFLT_URI
#     Reading alice
# Closing connection to DFLT_URI
#

with dao2:
    dao2.read('bob')
    dao2.read('alice')
print("")
# DAO entry...
# Opening connection to DFLT_URI
# Opening connection to DFLT_URI
#     Reading bob
# Closing connection to DFLT_URI
# Opening connection to DFLT_URI
#     Reading alice
# Closing connection to DFLT_URI
# ... DAO exit
# Closing connection to DFLT_URI
#

So really, we don't get much use out of our DAO's context manager here, yet.

What we'd like is for:

with dao2:
    dao2.read('bob')
    dao2.read('alice')
# DAO entry...
# Opening connection to DFLT_URI
#     Reading bob
#     Reading alice
# ... DAO exit
# Closing connection to DFLT_URI

to open the connection once, read bob and alice, and then close, once.

For this, we'll need the top level (DAO) context to signal to the lower context (the read) if the connection was already opened of not.
This would do the trick:

class DAO3:
    def __init__(self, uri='DFLT_URI'):
        self.connection = Connection(uri)
        self.connection_opened = False

    def read(self, k):
        if self.connection_opened:
            print(f"    Reading {k}")
        else:
            with self.connection:
                print(f"    Reading {k}")

    def __enter__(self):
        print(f"DAO entry...")
        self.connection.__enter__()
        self.connection_opened = True
        return self

    def __exit__(self, *exc_info):
        print(f"... DAO exit")
        self.connection_opened = False
        return self.connection.__exit__(*exc_info)

dao3 = DAO3()    
with dao3:
    dao3.read('bob')
    dao3.read('alice')
print("")
# DAO entry...
# Opening connection to DFLT_URI
#     Reading bob
#     Reading alice
# ... DAO exit
# Closing connection to DFLT_URI
#

Good, we got what we wanted, but it's not super clean, nor is it reusable (we need this for other methods too (write, delete, etc.).

To address the issue of cleanly separating concerns and making the "don't open/close twice" context manager concern reusable across different methods (like read, write, delete, etc.), we can introduce a more flexible connection context manager. This approach involves creating a separate context manager that can intelligently handle the state of the connection—whether it's already opened or not—without directly intertwining this logic with the DAO or connection classes themselves.

Approach

1. Flexible Connection Context Manager: Create a context manager that wraps the connection handling logic. This context manager would manage the state of the connection (open or closed) and prevent redundant opening or closing operations based on the current state.

2. Integration with DAO Classes: The DAO classes would utilize this flexible context manager to manage their connections, allowing them to focus on their primary responsibilities (like read, write, delete) without worrying about the connection state.

Implementation Steps

1. Define the Flexible Connection Context Manager: Implement a context manager class that accepts a connection object and manages opening and closing the connection based on whether it's already open. This could involve adding a flag to the connection class or managing state within the context manager itself.

2. Modify the Connection Class: Ensure the connection class supports idempotent open and close operations or at least doesn't fail when these operations are called redundantly.

3. Update DAO Classes to Use the Context Manager: Refactor the DAO classes to use the new context manager for connection management. This could mean modifying the __enter__ and __exit__ methods of the DAO classes to delegate connection management to the context manager.

Example Code

Here is a sketch of how the components might look:

class FlexibleConnectionManager:
    def __init__(self, connection):
        self.connection = connection
        self.owns_connection = False

    def __enter__(self):
        if not self.connection.is_open:
            self.connection.open()
            self.owns_connection = True
        return self.connection

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.owns_connection and self.connection.is_open:
            self.connection.close()
            self.owns_connection = False

class Connection:
    def __init__(self, uri='DFLT_URI'):
        self.uri = uri
        self.is_open = False

    def open(self):
        if not self.is_open:
            print(f"Opening connection to {self.uri}")
            self.is_open = True

    def close(self):
        if self.is_open:
            print(f"Closing connection to {self.uri}")
            self.is_open = False

    __enter__ = open
    __exit__ = close

class DAO:
    def __init__(self, uri='DFLT_URI'):
        self.connection = Connection(uri)

    def read(self, k):
        with FlexibleConnectionManager(self.connection):
            print(f"    Reading {k}")

    def __enter__(self):
        self.connection_manager = FlexibleConnectionManager(self.connection)
        return self.connection_manager.__enter__()

    def __exit__(self, *exc_info):
        self.connection_manager.__exit__(*exc_info)

Test it

When global entry:

dao = DAO()    
with dao:
    dao.read('bob')
    dao.read('alice')
# Opening connection to DFLT_URI
#     Reading bob
#     Reading alice
# Closing connection to DFLT_URI

Local entry:

dao.read('bob')
dao.read('alice')
# Opening connection to DFLT_URI
#     Reading bob
# Closing connection to DFLT_URI
# Opening connection to DFLT_URI
#     Reading alice
# Closing connection to DFLT_URI

Advantages of This Approach

• Separation of Concerns: This approach cleanly separates the connection management logic from the DAO's business logic, making the code more maintainable and reusable.
• Flexibility: The FlexibleConnectionManager can be used with any DAO class or method that requires connection management, reducing code duplication.
• Idempotency: The connection handling becomes idempotent, gracefully handling redundant open/close requests without error.

By encapsulating the connection state management in a dedicated context manager, you can simplify the DAO implementations and make the overall design more robust and flexible.