PostgreSQL 使用指南

使用用户名和密码进行身份验证

配方源码： postgresql_authenticate.py

要连接到 PostgreSQL 数据库，必须在 URI 中提供用户名和密码。例如，

postgresql://username:password@hostname:port/dbname

详细信息请参阅 PostgreSQL 文档。

import os

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
conn = adbc_driver_postgresql.dbapi.connect(uri)

with conn.cursor() as cur:
    cur.execute("SELECT 1")
    print(cur.fetchone())
    # Output: (1,)

conn.close()

标准输出

(1,)

启用自动提交模式

配方源码： postgresql_autocommit.py

您可以通过将 autocommit 参数传递给 connect 函数来启用自动提交模式。启用自动提交后，每条语句都会自动提交，无需显式调用 commit()。这对于无法在事务内运行的操作，或者当您希望立即提交每条语句时非常有用。

import os

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]

# Enable autocommit mode
conn = adbc_driver_postgresql.dbapi.connect(uri, autocommit=True)

with conn.cursor() as cur:
    # In autocommit mode, this statement is automatically committed
    cur.execute("CREATE TEMP TABLE IF NOT EXISTS autocommit_test (id INTEGER)")
    cur.execute("INSERT INTO autocommit_test VALUES (1)")

# Verify the data was committed
with conn.cursor() as cur:
    cur.execute("SELECT * FROM autocommit_test")
    assert cur.fetchone() == (1,)

conn.close()

从 Arrow 数据集创建/追加到表

配方源码： postgresql_create_dataset_table.py

ADBC 可以轻松地将 PyArrow 数据集加载到您的数据存储中。

import os
import tempfile
from pathlib import Path

import pyarrow
import pyarrow.csv
import pyarrow.dataset
import pyarrow.feather
import pyarrow.parquet

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
conn = adbc_driver_postgresql.dbapi.connect(uri)

为了进行测试，我们首先确保即将使用的表不存在。

with conn.cursor() as cur:
    cur.execute("DROP TABLE IF EXISTS csvtable")
    cur.execute("DROP TABLE IF EXISTS ipctable")
    cur.execute("DROP TABLE IF EXISTS pqtable")
    cur.execute("DROP TABLE IF EXISTS csvdataset")
    cur.execute("DROP TABLE IF EXISTS ipcdataset")
    cur.execute("DROP TABLE IF EXISTS pqdataset")

conn.commit()

生成示例数据

tempdir = tempfile.TemporaryDirectory(
    prefix="adbc-docs-",
    ignore_cleanup_errors=True,
)
root = Path(tempdir.name)
table = pyarrow.table(
    [
        [1, 1, 2],
        ["foo", "bar", "baz"],
    ],
    names=["ints", "strs"],
)

首先，我们将写入单个文件。

csv_file = root / "example.csv"
pyarrow.csv.write_csv(table, csv_file)

ipc_file = root / "example.arrow"
pyarrow.feather.write_feather(table, ipc_file)

parquet_file = root / "example.parquet"
pyarrow.parquet.write_table(table, parquet_file)

我们还将生成一些分区数据集。

csv_dataset = root / "csv_dataset"
pyarrow.dataset.write_dataset(
    table,
    csv_dataset,
    format="csv",
    partitioning=["ints"],
)

ipc_dataset = root / "ipc_dataset"
pyarrow.dataset.write_dataset(
    table,
    ipc_dataset,
    format="feather",
    partitioning=["ints"],
)

parquet_dataset = root / "parquet_dataset"
pyarrow.dataset.write_dataset(
    table,
    parquet_dataset,
    format="parquet",
    partitioning=["ints"],
)

将 CSV 文件加载到 PostgreSQL

我们可以直接将 pyarrow.RecordBatchReader（来自 open_csv）传递给 adbc_ingest。我们也可以传递 pyarrow.dataset.Dataset 或 pyarrow.dataset.Scanner。

with conn.cursor() as cur:
    reader = pyarrow.csv.open_csv(csv_file)
    cur.adbc_ingest("csvtable", reader, mode="create")

    reader = pyarrow.dataset.dataset(
        csv_dataset,
        format="csv",
        partitioning=["ints"],
    )
    cur.adbc_ingest("csvdataset", reader, mode="create")

conn.commit()

with conn.cursor() as cur:
    cur.execute("SELECT ints, strs FROM csvtable ORDER BY ints, strs ASC")
    assert cur.fetchall() == [(1, "bar"), (1, "foo"), (2, "baz")]

    cur.execute("SELECT ints, strs FROM csvdataset ORDER BY ints, strs ASC")
    assert cur.fetchall() == [(1, "bar"), (1, "foo"), (2, "baz")]

将 Arrow IPC (Feather) 文件加载到 PostgreSQL

with conn.cursor() as cur:
    reader = pyarrow.ipc.RecordBatchFileReader(ipc_file)

由于 PyArrow API 的特性，我们必须先将文件读入内存。

    cur.adbc_ingest("ipctable", reader.read_all(), mode="create")

不过，Dataset API 会将数据流式传输到内存中，然后再加载到 PostgreSQL 中。

    reader = pyarrow.dataset.dataset(
        ipc_dataset,
        format="feather",
        partitioning=["ints"],
    )
    cur.adbc_ingest("ipcdataset", reader, mode="create")

conn.commit()

with conn.cursor() as cur:
    cur.execute("SELECT ints, strs FROM ipctable ORDER BY ints, strs ASC")
    assert cur.fetchall() == [(1, "bar"), (1, "foo"), (2, "baz")]

    cur.execute("SELECT ints, strs FROM ipcdataset ORDER BY ints, strs ASC")
    assert cur.fetchall() == [(1, "bar"), (1, "foo"), (2, "baz")]

将 Parquet 文件加载到 PostgreSQL

with conn.cursor() as cur:
    reader = pyarrow.parquet.ParquetFile(parquet_file)
    cur.adbc_ingest("pqtable", reader.iter_batches(), mode="create")

    reader = pyarrow.dataset.dataset(
        parquet_dataset,
        format="parquet",
        partitioning=["ints"],
    )
    cur.adbc_ingest("pqdataset", reader, mode="create")

conn.commit()

with conn.cursor() as cur:
    cur.execute("SELECT ints, strs FROM pqtable ORDER BY ints, strs ASC")
    assert cur.fetchall() == [(1, "bar"), (1, "foo"), (2, "baz")]

    cur.execute("SELECT ints, strs FROM pqdataset ORDER BY ints, strs ASC")
    assert cur.fetchall() == [(1, "bar"), (1, "foo"), (2, "baz")]

清理

conn.close()
tempdir.cleanup()

从 Arrow 表创建/追加到表

配方源码： postgresql_create_append_table.py

ADBC 允许使用 Arrow 表来创建数据库表或向其追加数据。

import os

import pyarrow

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
conn = adbc_driver_postgresql.dbapi.connect(uri)

为了进行测试，我们首先确保即将使用的表不存在。

with conn.cursor() as cur:
    cur.execute("DROP TABLE IF EXISTS example")
    cur.execute("DROP TABLE IF EXISTS example2")

现在我们可以创建表了。

with conn.cursor() as cur:
    data = pyarrow.table(
        [
            [1, 2, None, 4],
        ],
        schema=pyarrow.schema(
            [
                ("ints", "int32"),
            ]
        ),
    )
    cur.adbc_ingest("example", data, mode="create")

conn.commit()

摄取后，我们可以获取结果。

with conn.cursor() as cur:
    cur.execute("SELECT * FROM example")
    assert cur.fetchone() == (1,)
    assert cur.fetchone() == (2,)

    cur.execute("SELECT COUNT(*) FROM example")
    assert cur.fetchone() == (4,)

如果我们尝试再次摄取，它会失败，因为表已经存在。

with conn.cursor() as cur:
    try:
        cur.adbc_ingest("example", data, mode="create")
    except conn.ProgrammingError:
        pass
    else:
        raise RuntimeError("Should have failed!")

conn.rollback()

相反，我们可以追加到表中。

with conn.cursor() as cur:
    cur.adbc_ingest("example", data, mode="append")

    cur.execute("SELECT COUNT(*) FROM example")
    assert cur.fetchone() == (8,)

我们还可以选择：如果表不存在则创建它，否则追加数据。

with conn.cursor() as cur:
    cur.adbc_ingest("example2", data, mode="create_append")

    cur.execute("SELECT COUNT(*) FROM example2")
    assert cur.fetchone() == (4,)

    cur.adbc_ingest("example2", data, mode="create_append")

    cur.execute("SELECT COUNT(*) FROM example2")
    assert cur.fetchone() == (8,)

最后，我们可以替换该表。

with conn.cursor() as cur:
    cur.adbc_ingest("example", data.slice(0, 2), mode="replace")

    cur.execute("SELECT COUNT(*) FROM example")
    assert cur.fetchone() == (2,)

conn.close()

创建/追加到临时表

配方源码： postgresql_create_temp_table.py

ADBC 也允许创建临时表并向其追加数据。

import os

import pyarrow

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
conn = adbc_driver_postgresql.dbapi.connect(uri)

为了进行测试，我们首先确保即将使用的表不存在。

with conn.cursor() as cur:
    cur.execute("DROP TABLE IF EXISTS example")

要创建临时表，只需指定 “temporary” 选项即可。

data = pyarrow.table(
    [
        [1, 2, None, 4],
    ],
    schema=pyarrow.schema(
        [
            ("ints", "int32"),
        ]
    ),
)

with conn.cursor() as cur:
    cur.adbc_ingest("example", data, mode="create", temporary=True)

conn.commit()

摄取后，我们可以获取结果。

with conn.cursor() as cur:
    cur.execute("SELECT * FROM example")
    assert cur.fetchone() == (1,)
    assert cur.fetchone() == (2,)

    cur.execute("SELECT COUNT(*) FROM example")
    assert cur.fetchone() == (4,)

临时表与常规表是分开的，即使它们具有相同的名称。

with conn.cursor() as cur:
    cur.adbc_ingest("example", data.slice(0, 2), mode="create", temporary=False)

conn.commit()

with conn.cursor() as cur:

因为我们有两个同名的表，所以必须在这里明确引用正常的临时表。

    cur.execute("SELECT COUNT(*) FROM public.example")
    assert cur.fetchone() == (2,)

    cur.execute("SELECT COUNT(*) FROM example")
    assert cur.fetchone() == (4,)

conn.close()

关闭连接后，临时表会被隐式删除。如果我们重新连接，该表将不存在；我们将只能看到“普通”表。

with adbc_driver_postgresql.dbapi.connect(uri) as conn:
    with conn.cursor() as cur:
        cur.execute("SELECT COUNT(*) FROM example")
        assert cur.fetchone() == (2,)

所有常规的摄取选项也适用于临时表。更多示例请参阅 从 Arrow 数据集创建/追加到表。

执行带有绑定参数的语句

配方源码： postgresql_execute_bind.py

ADBC 允许使用 Python 和 Arrow 值作为绑定参数。目前，PostgreSQL 驱动程序仅支持对不生成结果集的查询使用绑定参数。

import os

import pyarrow

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
conn = adbc_driver_postgresql.dbapi.connect(uri)

我们将创建一个示例表来测试。

with conn.cursor() as cur:
    cur.execute("DROP TABLE IF EXISTS example")
    cur.execute("CREATE TABLE example (ints INT, bigints BIGINT)")

conn.commit()

我们可以绑定 Python 值。

with conn.cursor() as cur:
    cur.executemany("INSERT INTO example VALUES ($1, $2)", [(1, 2), (3, 4)])

    cur.execute("SELECT SUM(ints) FROM example")
    assert cur.fetchone() == (4,)

注意

如果您习惯了类似 psycopg 等库用于绑定参数的格式字符串样式 %s 语法，请注意这不受支持——仅支持 PostgreSQL 原生的 $1 语法。

我们也可以绑定 Arrow 值。

with conn.cursor() as cur:
    data = pyarrow.record_batch(
        [
            [5, 6],
            [7, 8],
        ],
        names=["$1", "$2"],
    )
    cur.executemany("INSERT INTO example VALUES ($1, $2)", data)

    cur.execute("SELECT SUM(ints) FROM example")
    assert cur.fetchone() == (15,)

conn.close()

在不使用 COPY 的情况下执行语句

配方源码： postgresql_execute_nocopy.py

ADBC 驱动程序默认尝试使用 COPY 执行查询，因为这对于大数据集速度更快。但是，PostgreSQL 并不支持对所有类型的查询使用 COPY。例如，SHOW 查询将无法工作。在这种情况下，您可以显式禁用 COPY 优化。

import os

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
conn = adbc_driver_postgresql.dbapi.connect(uri)

该选项可以在创建游标时设置：

with conn.cursor(
    adbc_stmt_kwargs={
        adbc_driver_postgresql.StatementOptions.USE_COPY.value: False,
    }
) as cur:
    cur.execute("SHOW ALL")
    print(cur.fetch_arrow_table().schema)

或者也可以在事后设置：

with conn.cursor() as cur:
    cur.adbc_statement.set_options(
        **{
            adbc_driver_postgresql.StatementOptions.USE_COPY.value: False,
        }
    )
    cur.execute("SHOW ALL")
    print(cur.fetch_arrow_table().schema)

如果不使用该选项，查询将失败，因为驱动程序会尝试使用 COPY 执行查询。

with conn.cursor() as cur:
    try:
        cur.execute("SHOW ALL")
    except conn.Error:
        pass
    else:
        raise RuntimeError("Expected error")

conn.close()

标准输出 (stdout)

name: string
setting: string
description: string
name: string
setting: string
description: string

获取表的 Arrow 模式 (Schema)

配方源码： postgresql_get_table_schema.py

ADBC 允许您获取表的模式，并将其作为 Arrow 模式。

import os

import pyarrow

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
conn = adbc_driver_postgresql.dbapi.connect(uri)

我们将创建一些示例表进行测试。

with conn.cursor() as cur:
    cur.execute("DROP TABLE IF EXISTS example")
    cur.execute("CREATE TABLE example (ints INT, bigints BIGINT)")

    cur.execute("CREATE SCHEMA IF NOT EXISTS other_schema")
    cur.execute("DROP TABLE IF EXISTS other_schema.example")
    cur.execute("CREATE TABLE other_schema.example (strings TEXT, values INT)")

conn.commit()

默认情况下，假定为“活动”目录/模式。

assert conn.adbc_get_table_schema("example") == pyarrow.schema(
    [
        ("ints", "int32"),
        ("bigints", "int64"),
    ]
)

我们可以明确指定 PostgreSQL 模式，以获取不同命名空间中表的 Arrow 模式。

注意

在 PostgreSQL 中，您只能查询您连接到的数据库（目录）。因此我们不能在此处指定目录（或者说，这样做没有意义）。

请注意，NUMERIC 列被读取为字符串，因为 PostgreSQL 的小数类型无法直接映射到 Arrow 的小数类型。

assert conn.adbc_get_table_schema(
    "example",
    db_schema_filter="other_schema",
) == pyarrow.schema(
    [
        ("strings", "string"),
        ("values", "int32"),
    ]
)

conn.close()

获取查询的 Arrow 模式

配方源码： postgresql_get_query_schema.py

ADBC 允许您在不执行查询的情况下获取结果集的模式。

import os

import pyarrow

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
conn = adbc_driver_postgresql.dbapi.connect(uri)

我们将创建一个示例表来测试。

with conn.cursor() as cur:
    cur.execute("DROP TABLE IF EXISTS example")
    cur.execute("CREATE TABLE example (ints INT, bigints BIGINT)")

conn.commit()

expected = pyarrow.schema(
    [
        ("ints", "int32"),
        ("bigints", "int64"),
    ]
)

with conn.cursor() as cur:
    assert cur.adbc_execute_schema("SELECT * FROM example") == expected

PostgreSQL 在此处不知道类型，所以它只是返回一个猜测值。

    assert cur.adbc_execute_schema("SELECT $1 AS res") == pyarrow.schema(
        [
            ("res", "string"),
        ]
    )

conn.close()

列出目录、模式和表

配方源码： postgresql_list_catalogs.py

ADBC 允许列出数据库中的表、目录和模式。

import os

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
conn = adbc_driver_postgresql.dbapi.connect(uri)

我们将创建一个示例表来查找。

with conn.cursor() as cur:
    cur.execute("DROP TABLE IF EXISTS example")
    cur.execute("CREATE TABLE example (ints INT, bigints BIGINT)")

conn.commit()

数据以 PyArrow RecordBatchReader 的形式提供。

objects = conn.adbc_get_objects(depth="all").read_all()

为了方便起见，我们将把它转换为普通的 Python 数据。

objects = objects.to_pylist()
catalog = objects[0]
assert catalog["catalog_name"] == "postgres"

db_schema = catalog["catalog_db_schemas"][0]
assert db_schema["db_schema_name"] == "public"

tables = db_schema["db_schema_tables"]
example = [table for table in tables if table["table_name"] == "example"]
assert len(example) == 1
example = example[0]

assert example["table_columns"][0]["column_name"] == "ints"
assert example["table_columns"][1]["column_name"] == "bigints"

conn.close()

使用 SQLAlchemy 进行连接池管理

配方源码： postgresql_pool.py

ADBC 没有实现连接池，因为这通常不是 DBAPI 驱动程序的功能。相反，请使用第三方连接池，例如 SQLAlchemy 中内置的连接池。

import os

import sqlalchemy.pool

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]

source = adbc_driver_postgresql.dbapi.connect(uri)

adbc_driver_manager.dbapi.Connection.adbc_clone() 从现有连接打开一个新连接，尽可能共享内部资源。例如，PostgreSQL 驱动程序将共享内部 OID 缓存，从而节省连接开销。

pool = sqlalchemy.pool.QueuePool(source.adbc_clone, max_overflow=1, pool_size=2)

现在我们可以从池中获取连接；SQLAlchemy 会覆盖 close() 方法以将连接返回给池。

注意

与底层 ADBC 连接不同，SQLAlchemy 的包装器不支持上下文管理器协议。

conn = pool.connect()

assert pool.checkedin() == 0
assert pool.checkedout() == 1

with conn.cursor() as cur:
    cur.execute("SELECT 1")
    assert cur.fetchone() == (1,)

conn.close()

assert pool.checkedin() == 1
assert pool.checkedout() == 0

source.close()

结合使用 Pandas 和 ADBC

配方源码： postgresql_pandas.py

ADBC 集成到了流行的 dataframe 库 pandas 中。Pandas 可以使用 ADBC 与 PostgreSQL 及其他数据库交换数据。与使用 SQLAlchemy 或其他选项相比，将 ADBC 与 pandas 结合使用可以获得更好的性能，例如避免了与 Python 对象之间的过度转换。

import os

import pandas as pd

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
conn = adbc_driver_postgresql.dbapi.connect(uri)

我们将使用 pd.DataFrame.to_sql 创建一个示例表。

data = pd.DataFrame(
    {
        "ints": [1, 2, None, 4],
        "strs": ["a", "b", "c", "d"],
    }
)
data.to_sql("example", conn, if_exists="replace")
conn.commit()

创建表后，我们可以将 ADBC 连接和 SQL 查询传递给 pd.read_sql，以 pandas DataFrame 的形式获取结果集。

df = pd.read_sql("SELECT * FROM example WHERE ints > 1", conn)

assert len(df) == 2

conn.close()

与 ADBC 接口相比，pandas 提供了更方便、更高级的 API，特别是对于那些已经在使用 pandas 的用户而言。

结合使用 Polars 和 ADBC

配方源码： postgresql_polars.py

ADBC 可以与使用 Rust 编写的 dataframe 库 Polars 结合使用。根据其文档，

如果后端支持直接返回 Arrow 数据，则将使用此设施高效地实例化 DataFrame；否则，将从行数据初始化 DataFrame。

显然，ADBC 直接返回 Arrow 数据，这使得 ADBC 和 Polars 成为天作之合。

import os

import polars as pl

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]

我们将使用 Polars 和 polars.DataFrame.write_database() 来创建示例表。在使用 Polars 时，我们无需自己打开 ADBC 连接。

data = pl.DataFrame(
    {
        "ints": [1, 2, None, 4],
        "strs": ["a", "b", "c", "d"],
    }
)
data.write_database("example", uri, engine="adbc", if_table_exists="replace")

创建表后，我们可以使用 polars.read_database_uri() 来获取结果。同样，我们可以直接传入 URI 并告诉 Polars 为我们管理 ADBC。

df = pl.read_database_uri("SELECT * FROM example WHERE ints > 1", uri, engine="adbc")

assert len(df) == 2