dataset wrapper for With statement

[Daniil Okhlopkov]

I really love dataset library but it lacks a nice wrapper for with statement which will automatically close DB connection.

So I have to create my own wrapper and keep it as a submodule in my projects. It would be nice to have something like this built in the dataset future versions.

import dataset

class DB(object):
    """
        small wrapper around dataset which
        supports 'with' statements

        Example:
        with DB() as db:
            res = db.query("SELECT * FROM table LIMIT 10")
    """
    
    def __init__(self, database_url):
        self.database_url = database_url

    def __enter__(self):
        self.db = dataset.connect(self.database_url)
        return self.db
  
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.db.executable.close()

[Eugene Chernyavsky]

Can be better with schema:

class DB(object):

    def __init__(self, database_url: str = None, schema: str = None):
        if not database_url:
            database_url = os.environ.get("DATABASE_URL", "sqlite://")
        self.database_url = database_url.replace("postgres://", "postgresql://")
        self.connect_args = None
        self.schema = schema
        if schema:
            self.connect_args = {'connect_args': {'options': '-csearch_path={}'.format(schema)}}

    def __enter__(self):
        self.db = dataset.connect(self.database_url, schema=self.schema, engine_kwargs=self.connect_args)
        return self.db

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.db.commit()
        self.db.close()

[medecau]

Unfortunately the database object can already be used to enter a context block, but it is being used to start/end transactions.

dataset/dataset/database.py

Lines 155 to 169 in 5c2dc8d

 def __enter__(self): 
     """Start a transaction.""" 
     self.begin() 
     return self 
  
 def __exit__(self, error_type, error_value, traceback): 
     """End a transaction by committing or rolling back.""" 
     if error_type is None: 
         try: 
             self.commit() 
         except Exception: 
             with safe_reraise(): 
                 self.rollback() 
     else: 
         self.rollback() 

In hindsight the proposed behaviour makes more sense, with a transaction context still available from Database.transaction.

But changing this would break API.

@pudo would a PR with this behaviour have any chance of being merged?

[mijaba]

It's possible to get the same effect using just the standard library. Specifically, contextlib.closing will work, since dataset.Database has a close method:

import contextlib
import dataset
with contextlib.closing(dataset.connect()) as db:
    print(db.tables)

The close method isn't documented on Read the Docs, admittedly, but I can't imagine it is going to change.