如何使用 Polars 在 FastAPI 中模拟数据库连接？

Question

我想在使用

pyodbc.Connection

包读取数据库的 FastAPI 应用程序中模拟

polars

。这是初始的

main.py

文件：

from fastapi import Depends, FastAPI
import polars as pl
import pyodbc

import os

app = FastAPI()

def init_db() -> pyodbc.Connection:
    connstring = os.environ.get("AZURE_SQL_CONNECTIONSTRING", "no_env_var")
    return pyodbc.connect(connstring)


@app.get("/")
async def root(conn: pyodbc.Connection = Depends(init_db)) -> str:
    df = pl.read_database(query="SELECT * FROM test_table", connection=conn)
    return df.write_json()

Depends

类注入数据库连接以供稍后测试。我通过 pl.read_database 函数连接到数据库。我想创建一个单元测试用例而不连接到真实数据库。

如何在不连接真实数据库的情况下使用

pytest

测试 API 调用？

编辑：将

to_dicts

替换为

write_json

以提高转换效率。

Answer 1

首先可以使用初始化数据库时可以返回的

polars.ConnectionOrCursor

类：

from fastapi import Depends, FastAPI
import polars as pl
from polars.type_aliases import ConnectionOrCursor
import pyodbc

import os

app = FastAPI()

def init_db() -> ConnectionOrCursor:
    connstring = os.environ.get("AZURE_SQL_CONNECTIONSTRING", "no_env_var")
    return pyodbc.connect(connstring)


@app.get("/")
async def root(conn: ConnectionOrCursor = Depends(init_db)) -> str:
    df = pl.read_database(query="SELECT * FROM test_table", connection=conn)
    return df.write_json()

之后，就可以在单元测试中使用假数据库了。您可以使用内存 SQLite 数据库。此外，

TestClient

类可用于测试FastAPI客户端。有关更多信息，请访问测试 FastAPI 文档。

这是一个示例测试文件：

from fastapi.testclient import TestClient
from polars.type_aliases import ConnectionOrCursor

import sqlite3

from your.source.module import app, init_db

persons = [
    (1, "John", 20),
    (2, "Doe", 30),
]


def setup_database() -> ConnectionOrCursor:
    connection = sqlite3.connect(":memory:", check_same_thread=False)
    cursor = connection.cursor()
    cursor.execute(
        "CREATE TABLE test_table(id, name, age)"
    )
    for person in persons:
        cursor.execute("INSERT INTO test_table VALUES(?,?,?,?)", person)
    connection.commit()

    return connection


client = TestClient(app)
app.dependency_overrides[init_db] = setup_database


def test_root():
    response = client.get("/")
    assert response.status_code == 200

如果有多个测试用例，请在连接构造函数中使用

check_same_thread=False

。可以用数据填充 SQLite 数据库。由于

sqlite.Connection

与 Polars 兼容，因此您可以覆盖

init_db

依赖项。

app.dependency_overrides

属性字典会覆盖 FastAPI 依赖项。

注意：真假数据库的表名（

test_table

）必须相同。