Last active
July 8, 2024 12:46
-
-
Save danielrichman/6046307 to your computer and use it in GitHub Desktop.
nicer postgres connection class & flask postgres
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| from __future__ import unicode_literals | |
| import logging | |
| import flask | |
| import psycopg2 | |
| import psycopg2.extras | |
| class PostgreSQL(object): | |
| """ | |
| A PostgreSQL helper extension for Flask apps | |
| On initialisation it adds an after_request function that commits the | |
| transaction (so that if the transaction rolls back the request will | |
| fail) and a app context teardown function that disconnects any active | |
| connection. | |
| You can of course (and indeed should) use :meth:`commit` if you need to | |
| ensure some changes have made it to the database before performing | |
| some other action. :meth:`teardown` is also available to be called | |
| directly. | |
| Connections are created by ``psycopg2.connenct(app.config["POSTGRES"])``. | |
| """ | |
| def __init__(self, app=None): | |
| self.app = app | |
| self.logger = logging.getLogger(__module__ + ".PostgreSQL") | |
| if app is not None: | |
| self.init_app(app) | |
| def init_app(self, app): | |
| """ | |
| Initialises the app by adding hooks | |
| * Hook: ``app.after_request(self.commit)`` | |
| * Hook: ``app.teardown_appcontext(self.teardown)`` | |
| """ | |
| app.after_request(self.commit) | |
| app.teardown_appcontext(self.teardown) | |
| def _connect(self): | |
| """Returns (always) a new connection to the database""" | |
| s = flask.current_app.config["POSTGRES"] | |
| self.logger.debug("connecting (%r)", s) | |
| return psycopg2.connect(s) | |
| @property | |
| def connection(self): | |
| """ | |
| Gets the PostgreSQL connection for this Flask request | |
| If no connection has been used in this request, it connects to the | |
| database. Further use of this property will reference the same | |
| connection | |
| The connection is committed and closed at the end of the request. | |
| """ | |
| g = flask.current_app.g | |
| if not hasattr(g, '_postgresql'): | |
| g._postgresql = self._connect() | |
| return g._postgresql | |
| def cursor(self, real_dict_cursor=False): | |
| """ | |
| Get a new postgres cursor for immediate use during a request | |
| If a cursor has not yet been used in this request, it connects to the | |
| database. Further cursors re-use the per-request connection. | |
| The connection is committed and closed at the end of the request. | |
| If real_dict_cursor is set, a RealDictCursor is returned | |
| """ | |
| kwargs = {} | |
| if real_dict_cursor: | |
| kwargs["cursor_factory"] = psycopg2.extras.RealDictCursor | |
| return self.connection.cursor(**kwargs) | |
| def commit(self): | |
| """ | |
| (Almost an) alias for self.connection.commit() | |
| ... except if self.connection has never been used (i.e., it does | |
| nothing this is a noop) | |
| """ | |
| g = flask.current_app.g | |
| if hasattr(g, '_postgresql'): | |
| self.logger.debug("committing") | |
| g._postgresql.commit() | |
| def teardown(self, exception): | |
| """Tears down any connection on current_app.g""" | |
| g = flask.current_app.g | |
| if hasattr(g, '_postgresql'): | |
| self.logger("teardown: closing") | |
| g._postgresql.close() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment