Psycopg2 和 Flask - 将连接绑定到 before_request 和 teardown_appcontext

Question

干杯，伙计们， 重构我的 Flask 应用程序时，我被困在将 db 连接绑定到 @app.before_request 并在 @app.teardown_appcontext 处关闭它。我正在使用普通的 Psycopg2 和应用工厂模式。

首先我创建了一个在应用工厂中调用的函数，这样我就可以将@app 用作suggested by Miguel Grinberg here：

def create_app(test_config=None):
    app = Flask(__name__, instance_relative_config=True)

    --

    from shop.db import connect_and_close_db
    connect_and_close_db(app)

    --

    return app

然后我尝试了http://flask.pocoo.org/docs/1.0/appcontext/#storing-data建议的这种模式：

def connect_and_close_db(app):

    @app.before_request
    def get_db_test():
        conn_string = "dbname=testdb user=testuser password=test host=localhost"
        if 'db' not in g:
            g.db = psycopg2.connect(conn_string)
        return g.db

    @app.teardown_appcontext
    def close_connection(exception):
        db = g.pop('db', None)

        if db is not None:
            db.close()

结果：

TypeError: 'psycopg2.extensions.connection' object is not callable

有人知道发生了什么以及如何使其发挥作用吗？

此外，我想知道一旦创建游标绑定到before_request，我将如何访问连接对象以创建游标？

Answer 1

这个解决方案可能远非完美，而且它并不是真正的 DRY。我欢迎 cmets 或基于此的其他答案。

要实现原始psycopg2 支持，您可能需要查看connection pooler。还有一个good guide 介绍如何使用 Flask 实现这一点。

基本思想是首先创建连接池。您希望在烧瓶应用程序初始化时建立它（这可以在 python 解释器中或通过可能有几个的 gunicorn 工作者 - 在这种情况下，每个工作者都有自己的连接池）。我选择将返回的池存储在配置中：

from flask import Flask, g, jsonify

import psycopg2
from psycopg2 import pool

app = Flask(__name__)

app.config['postgreSQL_pool'] = psycopg2.pool.SimpleConnectionPool(1, 20,
    user = "postgres",
    password = "very_secret",
    host = "127.0.0.1",
    port = "5432",
    database = "postgres")

注意SimpleConnectionPool 的前两个参数是min 和max 连接。这就是连接到您的数据库服务器的连接数，在这种情况下是1 和20。

接下来定义一个get_db函数：

def get_db():
    if 'db' not in g:
        g.db = app.config['postgreSQL_pool'].getconn()
    return g.db

这里使用的SimpleConnectionPool.getconn() 方法只是从池中返回一个连接，我们将其分配给g.db 并返回。这意味着当我们在代码中的任何地方调用get_db() 时，它会返回相同的连接，或者如果不存在则创建一个连接。不需要 before.context 装饰器。

定义你的拆解函数：

@app.teardown_appcontext
def close_conn(e):
    db = g.pop('db', None)
    if db is not None:
        app.config['postgreSQL_pool'].putconn(db)

这在应用上下文被销毁时运行，并使用SimpleConnectionPool.putconn() 收起连接。

最后定义一条路线：

@app.route('/')
def index():
    db = get_db()
    cursor = db.cursor()

    cursor.execute("select 1;")
    result = cursor.fetchall()
    print (result)

    cursor.close()
    return jsonify(result)

此代码适用于我针对在 docker 容器中运行的 postgres 进行的测试。一些可能需要改进的地方：

• 这个视图不是很干。也许您可以将其中的一些内容移到get_db 函数中，以便它返回一个游标。 (!!!)

• 当python解释器退出时，你也应该找个地方关闭与app.config['postgreSQL_pool'].closeall的连接

• 虽然测试了某种监控池的方法会很好，但您可以在负载下观察池/数据库连接并确保池器按预期运行。

(!!!)在另一个地方，sqlalchemy.scoped_sessiondocumentation 解释了更多与此相关的事情，并提供了一些关于其“会话”如何与请求相关的理论。他们以这样一种方式实现了它，您可以调用Session.query('SELECT 1')，如果它不存在，它将创建会话。

---

编辑：这是一个gist，其中包含您的应用工厂模式，以及评论中的示例用法。

Answer 2

目前我正在使用这种模式： （如果我想出更好的解决方案，我最终会编辑这个答案）

这是我们使用数据库的主要脚本。它使用配置中的两个函数：get_db() 从池中获取连接，put_db() 将连接返回到池中：

from config import get_db, put_db
from threading import Thread
from time import sleep

def select():
    db = get_db()
    sleep(1)
    cursor = db.cursor()
    # Print select result and db connection address in memory
    # To see if it gets connection from another addreess on second thread
    cursor.execute("SELECT 'It works %s'", (id(db),))
    print(cursor.fetchone())
    cursor.close()
    put_db(db)

Thread(target=select).start()
Thread(target=select).start()
print('Main thread')

这是config.py：

import sys
import os
import psycopg2
from psycopg2 import pool
from dotenv import load_dotenv, find_dotenv
load_dotenv(find_dotenv())

def get_db(key=None):
    return getattr(get_db, 'pool').getconn(key)

def put_db(conn, key=None):
    getattr(get_db, 'pool').putconn(conn, key=key)

# So we here need to init connection pool in main thread in order everything to work
# Pool is initialized under function object get_db
try:
    setattr(get_db, 'pool', psycopg2.pool.ThreadedConnectionPool(1, 20, os.getenv("DB")))
    print(color.red('Initialized db'))
except psycopg2.OperationalError as e:
    print(e)
    sys.exit(0)

如果你好奇的话，还有一个 .env 文件，其中包含 DB 环境变量中的数据库连接字符串：

DB="dbname=postgres user=postgres password=1234 host=127.0.0.1 port=5433"

（.env 文件是使用config.py 中的dotenv 模块加载的）