Introducing a good way to manage DB connections in Python

Introduction

Looking at a certain OSS code, there was something that looked good, so let me introduce it.

code

A certain OSS is Apache Airflow. Among them, airflow / utils / session.py was a good feeling.

easy explanation

First, from session.py.

import contextlib
from functools import wraps

from airflow import settings

# contextlib.If you specify contextmanager, it will automatically close using with.
@contextlib.contextmanager
def create_session():
    """
    Contextmanager that will create and teardown a session.
    """
    session = settings.Session()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()

#I want to use Session It complements Session nicely when used in a function
def provide_session(func):
    """
    Function decorator that provides a session if it isn't provided.
    If you want to reuse a session or run the function as part of a
    database transaction, you pass it to the function, if not this wrapper
    will create one and close it for you.
    """
    @wraps(func)
    def wrapper(*args, **kwargs):
        arg_session = 'session'

        func_params = func.__code__.co_varnames
        session_in_args = arg_session in func_params and \
            func_params.index(arg_session) < len(args)
        session_in_kwargs = arg_session in kwargs
        #If there is a session in the argument of function, use it, if not, create it
        if session_in_kwargs or session_in_args:
            return func(*args, **kwargs)
        else:
            with create_session() as session:
                kwargs[arg_session] = session
                return func(*args, **kwargs)

    return wrapper

So, the code on the side to use next. Line 940 in airflow / models / baseoperator.py /airflow/blob/master/airflow/models/baseoperator.py#L940)

@provide_session
def get_task_instances(self, start_date: Optional[datetime] = None,
                        end_date: Optional[datetime] = None,
                        session: Session = None) -> List[TaskInstance]:
    """
    Get a set of task instance related to this task for a specific date
    range.
    """
    end_date = end_date or timezone.utcnow()
    return session.query(TaskInstance)\
        .filter(TaskInstance.dag_id == self.dag_id)\
        .filter(TaskInstance.task_id == self.task_id)\
        .filter(TaskInstance.execution_date >= start_date)\
        .filter(TaskInstance.execution_date <= end_date)\
        .order_by(TaskInstance.execution_date)\
        .all()

If you specify it as a decorator with a function that uses session like If session is not set in the caller, it will be newly created (and will be closed at the end). If you have set session, you can achieve a nice feeling of using it as it is.

At the end

I think this method can be applied in various ways other than DB connection.

reference

Recommended Posts

Introducing a good way to manage DB connections in Python
A clever way to time processing in Python
[python] Manage functions in a list
A simple way to avoid multiple for loops in Python
A standard way to develop and distribute packages in Python
How to get a stacktrace in python
Easy way to use Wikipedia in Python
Manage python packages to install in containers
A way to understand Python duck typing
What is the fastest way to create a reverse dictionary in python?
Try to calculate a statistical problem in Python
How to clear tuples in a list (Python)
To execute a Python enumerate function in JavaScript
How to embed a variable in a python string
I want to create a window in Python
How to create a JSON file in Python
To add a module to python put in Julialang
How to notify a Discord channel in Python
[Python] How to draw a histogram in Matplotlib
Tips for using ElasticSearch in a good way
Searching for an efficient way to write a Dockerfile in Python with poetry
Introducing Python 2.7 to CentOS 6.6
Parse a JSON string written to a file in Python
I want to easily implement a timeout in python
Try to make a Python module in C language
I want to write in Python! (2) Let's write a test
Create a plugin to run Python Doctest in Vim (2)
I tried to implement a pseudo pachislot in Python
A memorandum to run a python script in a bat file
I want to randomly sample a file in Python
I want to work with a robot in python.
Things to note when initializing a list in Python
[Python] Created a method to convert radix in 1 second
How to execute a command using subprocess in Python
Fixed a way to force Windows to boot in UEFI
Publish / upload a library created in Python to PyPI
[Mac] A super-easy way to execute system commands in Python and output the results
Create a function in Python
To flush stdout in Python
Create a dictionary in Python
I tried to implement a one-dimensional cellular automaton in Python
How to slice a block multiple array from a multiple array in Python
A road to intermediate Python
A story about how to specify a relative path in python.
Change the standard output destination to a file in Python
Login to website in Python
Probably the easiest way to create a pdf with Python3
How to import a file anywhere you like in Python
Developed a library to get Kindle collection list in Python
How to define multiple variables in a python for statement
I tried "How to get a method decorated in Python"
How to develop in a virtual environment of Python [Memo]
Make a bookmarklet in Python
To return char * in a callback function using ctypes in Python
How to get the last (last) value in a list in Python
How to get a list of built-in exceptions in python
[python] Manage functions in a dictionary (command table, function table, function pointer)
Introducing Python in Practice (PiP)
[Python] Another way to import
Create a tool to check scraping rules (robots.txt) in Python
I tried to make a stopwatch using tkinter in python