From 476adbb4191c80dc2ec203a74a487b271d733df7 Mon Sep 17 00:00:00 2001 From: paomian Date: Thu, 29 Oct 2026 20:13:48 +0800 Subject: [PATCH 1/2] Refactor GreptimeDB Python SDK to separate gRPC and PostgreSQL backends - Updated documentation to clarify backend usage and import paths. - Removed automatic backend selection; users must explicitly import the desired backend. - Simplified error handling by consolidating GreptimeError into a single implementation. - Enhanced type handling in the gRPC and PostgreSQL modules to ensure compatibility with shared schema. - Removed legacy native module and associated imports. - Added tests to verify that both backends correctly utilize the shared schema and handle errors appropriately. - Ensured that the public schema is consistent across both backends. --- AGENTS.md | 2 +- README.md | 7 +- README.zh-CN.md | 7 +- docs/usage.md | 11 +- examples/bench_write.py | 3 +- examples/bulk_example.py | 2 +- examples/insert_example.py | 2 +- packages/grpc/src/bulk.rs | 7 +- packages/grpc/src/error.rs | 20 +- packages/grpc/src/lib.rs | 4 - packages/grpc/src/types.rs | 252 +++++++------------------ python/greptimedb_ingester/__init__.py | 37 +--- python/greptimedb_ingester/_errors.py | 15 +- python/greptimedb_ingester/_native.py | 8 - python/greptimedb_ingester/_types.py | 6 +- python/greptimedb_ingester/grpc.py | 8 - python/greptimedb_ingester/postgres.py | 118 ++++-------- tests/test_backends.py | 37 ++-- tests/test_convert.py | 12 +- tests/test_integration.py | 5 +- tests/test_postgres.py | 12 +- tests/test_postgres_schema_compat.py | 41 ---- tests/test_shared_schema.py | 49 +++++ tests/test_types.py | 12 +- 24 files changed, 230 insertions(+), 447 deletions(-) delete mode 100644 python/greptimedb_ingester/_native.py delete mode 100644 tests/test_postgres_schema_compat.py create mode 100644 tests/test_shared_schema.py diff --git a/AGENTS.md b/AGENTS.md index ec1d3f7..76ef8cd 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -19,7 +19,7 @@ Python write client for GreptimeDB. Package name `greptimedb-ingester`, import ` ## Build and release - Root `uv build` builds the pure Python facade. Build the native dependency from `packages/grpc` with maturin. Keep facade/plugin versions, the exact grpc extra pin and lock files synchronized with `scripts/set-version.py`. -- Test base, postgres-only, gRPC-only and combined installed wheels. Default root Client prefers installed gRPC, otherwise postgres; explicit `.grpc.Client` and `.postgres.Client` select a backend. +- Test base, postgres-only, gRPC-only and combined installed wheels. The root exports only shared schema and GreptimeError. Import clients explicitly from `.grpc` or `.postgres`; do not add automatic backend selection or legacy import shims. - `v*` tags publish the native dependency first, then the public facade, using `PYPI_TOKEN` with access to both distributions. Both backend integrations must pass. - Wheels are built on the target OS: Linux x86_64, Linux aarch64, Windows x64, macOS Apple Silicon. Do not cross-compile those wheels. diff --git a/README.md b/README.md index 57821fb..f266c2b 100644 --- a/README.md +++ b/README.md @@ -35,7 +35,7 @@ print(affected) `insert` creates missing tables and writes COPY batches (default: 10,000 rows). Batches commit separately. The server must support `COPY FROM STDIN`. -See the [API reference](docs/api.md#postgresql-copy-后端) for supported types and timestamp rules. +See the [COPY example](examples/postgres_example.py). ## gRPC @@ -47,8 +47,7 @@ client = Client(["127.0.0.1:4001"], database="public") ``` Supports `insert`, `delete`, and `bulk_writer`. See [usage](docs/usage.md) and [examples](examples). -With both backends installed, the root `Client` defaults to gRPC; use the explicit -imports above to select a backend. +Import `Client` from the backend you use. Import shared schema types from `greptimedb_ingester`. ## Development @@ -58,4 +57,4 @@ uv run --extra postgres pytest ``` For both backends, use `--all-extras` instead of `--extra postgres` (requires Rust 1.85+). -See [testing and release](docs/testing.md) for integration tests and publishing. +See the [CI workflow](.github/workflows/wheels.yml) for integration tests and publishing. diff --git a/README.zh-CN.md b/README.zh-CN.md index 126b67f..4cc98e2 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -35,7 +35,7 @@ print(affected) `insert` 自动创建不存在的表并分批 COPY，默认每批 10,000 行。 各批独立提交。服务端须支持 `COPY FROM STDIN`。 -支持类型和时间戳规则见 [API](docs/api.md#postgresql-copy-后端)。 +完整用法见 [COPY 示例](examples/postgres_example.py)。 ## gRPC @@ -47,8 +47,7 @@ client = Client(["127.0.0.1:4001"], database="public") ``` 支持 `insert`、`delete` 和 `bulk_writer`，详见 [使用指南](docs/usage.md) 和 [示例](examples)。 -两种后端都安装时，顶层 `Client` 默认 gRPC； -使用上面的显式导入选择后端。 +从对应后端导入 `Client`，公共 schema 类型从 `greptimedb_ingester` 导入。 ## 开发 @@ -58,4 +57,4 @@ uv run --extra postgres pytest ``` 开发两种后端时，将 `--extra postgres` 换成 `--all-extras`，需要 Rust 1.85+。 -集成测试与发布流程见 [测试文档](docs/testing.md)。 +集成测试与发布流程见 [CI 配置](.github/workflows/wheels.yml)。 diff --git a/docs/usage.md b/docs/usage.md index 623a533..3198be5 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -29,7 +29,7 @@ uv sync --python 3.10 --extra grpc ## 连接 ```python -from greptimedb_ingester import Client +from greptimedb_ingester.grpc import Client client = Client( ["127.0.0.1:4001"], @@ -57,7 +57,8 @@ client.health_check() ```python from datetime import datetime, timezone -from greptimedb_ingester import Client, Column, ColumnDataType, SemanticType +from greptimedb_ingester import Column, ColumnDataType, SemanticType +from greptimedb_ingester.grpc import Client client = Client(["127.0.0.1:4001"], database="public") @@ -111,7 +112,8 @@ client.delete( ```python from datetime import datetime, timezone -from greptimedb_ingester import Client, Column, ColumnDataType, SemanticType, WriteOptions +from greptimedb_ingester import Column, ColumnDataType, SemanticType +from greptimedb_ingester.grpc import Client, WriteOptions client = Client(["127.0.0.1:4001"], database="public") columns = [ @@ -260,8 +262,7 @@ with Client(["127.0.0.1:4003"], batch_size=5000) as client: affected = client.insert("copy_data", columns, rows) ``` -只安装 `[postgres]` 时，也可直接导入顶层 `Client`。两种后端都安装时， -顶层 `Client` 默认 gRPC，COPY 用显式导入；schema 和 `GreptimeError` 共用。 +两个后端分别从 `greptimedb_ingester.grpc`、`greptimedb_ingester.postgres` 导入 `Client`；公共 schema 和 `GreptimeError` 从顶层导入。 COPY 只提供 insert、health_check、close 和 with；不提供 bulk_writer、delete 或 gRPC hints。 支持类型、时间转换、逐批提交和失败行为见 [README](../README.zh-CN.md) 和 [API](api.md#postgresql-copy-后端)。完整例子见 `examples/postgres_example.py`。 diff --git a/examples/bench_write.py b/examples/bench_write.py index 315ae87..9e2f167 100644 --- a/examples/bench_write.py +++ b/examples/bench_write.py @@ -7,7 +7,8 @@ import argparse import time -from greptimedb_ingester import Client, Column, ColumnDataType, SemanticType, WriteOptions +from greptimedb_ingester import Column, ColumnDataType, SemanticType +from greptimedb_ingester.grpc import Client, WriteOptions COLUMNS = [ Column("device_id", ColumnDataType.STRING, SemanticType.TAG), diff --git a/examples/bulk_example.py b/examples/bulk_example.py index de55762..fdec363 100644 --- a/examples/bulk_example.py +++ b/examples/bulk_example.py @@ -7,7 +7,7 @@ from datetime import timedelta -from greptimedb_ingester import Client, WriteOptions +from greptimedb_ingester.grpc import Client, WriteOptions from ddl import drop_table from type_samples import BULK_COLUMNS, BULK_ROW, INSERT_ROW, WHEN diff --git a/examples/insert_example.py b/examples/insert_example.py index 0e07ea6..1d96569 100644 --- a/examples/insert_example.py +++ b/examples/insert_example.py @@ -8,7 +8,7 @@ from datetime import timedelta -from greptimedb_ingester import Client +from greptimedb_ingester.grpc import Client from ddl import drop_table from type_samples import INSERT_COLUMNS, INSERT_ROW, WHEN diff --git a/packages/grpc/src/bulk.rs b/packages/grpc/src/bulk.rs index adad700..abb2c18 100644 --- a/packages/grpc/src/bulk.rs +++ b/packages/grpc/src/bulk.rs @@ -135,8 +135,11 @@ impl BulkWriter { } #[getter] - fn columns(&self) -> Vec { - self.columns.clone() + fn columns(&self, py: Python<'_>) -> PyResult>> { + self.columns + .iter() + .map(|column| column.to_python(py)) + .collect() } fn alloc_rows(&self, py: Python<'_>, capacity: usize) -> PyResult> { diff --git a/packages/grpc/src/error.rs b/packages/grpc/src/error.rs index 7b00437..b73911f 100644 --- a/packages/grpc/src/error.rs +++ b/packages/grpc/src/error.rs @@ -1,13 +1,5 @@ -use pyo3::exceptions::PyException; use pyo3::prelude::*; -pyo3::create_exception!( - greptimedb_ingester, - GreptimeError, - PyException, - "Error returned by the GreptimeDB ingester." -); - #[derive(Debug)] pub struct RemoteError { pub message: String, @@ -36,12 +28,14 @@ impl From for RemoteError { pub fn raise(message: impl Into, retriable: bool) -> PyErr { let message = message.into(); Python::with_gil(|py| { - let err = GreptimeError::new_err(message.clone()); - let instance = err.value(py); - if instance.setattr("retriable", retriable).is_err() { - return GreptimeError::new_err(format!("retriable={retriable}: {message}")); + match py + .import("greptimedb_ingester._errors") + .and_then(|module| module.getattr("GreptimeError")) + .and_then(|class| class.call1((message, retriable))) + { + Ok(instance) => PyErr::from_value(instance), + Err(error) => error, } - err }) } diff --git a/packages/grpc/src/lib.rs b/packages/grpc/src/lib.rs index ee464d9..f54a751 100644 --- a/packages/grpc/src/lib.rs +++ b/packages/grpc/src/lib.rs @@ -22,10 +22,6 @@ pub fn runtime() -> &'static tokio::runtime::Runtime { #[pymodule] fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> { - m.add("GreptimeError", m.py().get_type::())?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/packages/grpc/src/types.rs b/packages/grpc/src/types.rs index e5adf0d..d6c2d8c 100644 --- a/packages/grpc/src/types.rs +++ b/packages/grpc/src/types.rs @@ -6,113 +6,41 @@ use crate::error::raise; use greptimedb_ingester::api::v1::ColumnDataType as ProtoType; use greptimedb_ingester::{CompressionType, GrpcCompression}; -#[pyclass(eq, eq_int, frozen, name = "SemanticType")] #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum SemanticType { - #[pyo3(name = "TAG")] Tag = 0, - #[pyo3(name = "FIELD")] Field = 1, - #[pyo3(name = "TIMESTAMP")] Timestamp = 2, } -#[pymethods] -impl SemanticType { - fn __repr__(&self) -> String { - format!("SemanticType.{self:?}") - .replace("Tag", "TAG") - .replace("Field", "FIELD") - .replace("Timestamp", "TIMESTAMP") - } -} - -#[pyclass(eq, eq_int, frozen, name = "ColumnDataType")] #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum ColumnDataType { - #[pyo3(name = "BOOLEAN")] Boolean = 0, - #[pyo3(name = "INT8")] Int8 = 1, - #[pyo3(name = "INT16")] Int16 = 2, - #[pyo3(name = "INT32")] Int32 = 3, - #[pyo3(name = "INT64")] Int64 = 4, - #[pyo3(name = "UINT8")] Uint8 = 5, - #[pyo3(name = "UINT16")] Uint16 = 6, - #[pyo3(name = "UINT32")] Uint32 = 7, - #[pyo3(name = "UINT64")] Uint64 = 8, - #[pyo3(name = "FLOAT32")] Float32 = 9, - #[pyo3(name = "FLOAT64")] Float64 = 10, - #[pyo3(name = "BINARY")] Binary = 11, - #[pyo3(name = "STRING")] String = 12, - #[pyo3(name = "DATE")] Date = 13, - #[pyo3(name = "TIMESTAMP_SECOND")] TimestampSecond = 15, - #[pyo3(name = "TIMESTAMP_MILLISECOND")] TimestampMillisecond = 16, - #[pyo3(name = "TIMESTAMP_MICROSECOND")] TimestampMicrosecond = 17, - #[pyo3(name = "TIMESTAMP_NANOSECOND")] TimestampNanosecond = 18, - #[pyo3(name = "TIME_SECOND")] TimeSecond = 19, - #[pyo3(name = "TIME_MILLISECOND")] TimeMillisecond = 20, - #[pyo3(name = "TIME_MICROSECOND")] TimeMicrosecond = 21, - #[pyo3(name = "TIME_NANOSECOND")] TimeNanosecond = 22, - #[pyo3(name = "DECIMAL128")] Decimal128 = 30, - #[pyo3(name = "JSON")] Json = 31, } -#[pymethods] -impl ColumnDataType { - fn __repr__(&self) -> String { - let name = match self { - Self::Boolean => "BOOLEAN", - Self::Int8 => "INT8", - Self::Int16 => "INT16", - Self::Int32 => "INT32", - Self::Int64 => "INT64", - Self::Uint8 => "UINT8", - Self::Uint16 => "UINT16", - Self::Uint32 => "UINT32", - Self::Uint64 => "UINT64", - Self::Float32 => "FLOAT32", - Self::Float64 => "FLOAT64", - Self::Binary => "BINARY", - Self::String => "STRING", - Self::Date => "DATE", - Self::TimestampSecond => "TIMESTAMP_SECOND", - Self::TimestampMillisecond => "TIMESTAMP_MILLISECOND", - Self::TimestampMicrosecond => "TIMESTAMP_MICROSECOND", - Self::TimestampNanosecond => "TIMESTAMP_NANOSECOND", - Self::TimeSecond => "TIME_SECOND", - Self::TimeMillisecond => "TIME_MILLISECOND", - Self::TimeMicrosecond => "TIME_MICROSECOND", - Self::TimeNanosecond => "TIME_NANOSECOND", - Self::Decimal128 => "DECIMAL128", - Self::Json => "JSON", - }; - format!("ColumnDataType.{name}") - } -} - impl ColumnDataType { pub fn proto(self) -> ProtoType { match self { @@ -142,38 +70,8 @@ impl ColumnDataType { Self::Json => ProtoType::Json, } } - - pub fn name(self) -> &'static str { - match self { - Self::Boolean => "boolean", - Self::Int8 => "int8", - Self::Int16 => "int16", - Self::Int32 => "int32", - Self::Int64 => "int64", - Self::Uint8 => "uint8", - Self::Uint16 => "uint16", - Self::Uint32 => "uint32", - Self::Uint64 => "uint64", - Self::Float32 => "float32", - Self::Float64 => "float64", - Self::Binary => "binary", - Self::String => "string", - Self::Date => "date", - Self::TimestampSecond => "timestamp_second", - Self::TimestampMillisecond => "timestamp_millisecond", - Self::TimestampMicrosecond => "timestamp_microsecond", - Self::TimestampNanosecond => "timestamp_nanosecond", - Self::TimeSecond => "time_second", - Self::TimeMillisecond => "time_millisecond", - Self::TimeMicrosecond => "time_microsecond", - Self::TimeNanosecond => "time_nanosecond", - Self::Decimal128 => "decimal128", - Self::Json => "json", - } - } } -#[pyclass(name = "Column")] #[derive(Clone, Debug, PartialEq, Eq)] pub struct Column { pub(crate) name: String, @@ -183,100 +81,76 @@ pub struct Column { pub(crate) scale: Option, } -#[pymethods] -impl Column { - #[new] - #[pyo3(signature = (name, data_type, semantic_type, precision=None, scale=None))] - fn new( - name: String, - data_type: ColumnDataType, - semantic_type: SemanticType, - precision: Option, - scale: Option, - ) -> PyResult { - if name.is_empty() { - return Err(raise("column name must not be empty", false)); - } - match (data_type, precision, scale) { - (ColumnDataType::Decimal128, Some(precision), Some(scale)) => { - if precision == 0 || precision > 38 { - return Err(raise( - "decimal128 precision must be between 1 and 38", - false, - )); - } - if scale < 0 || i32::from(scale) > i32::from(precision) { - return Err(raise( - "decimal128 scale must be between 0 and precision", - false, - )); - } - } - (ColumnDataType::Decimal128, _, _) => { - return Err(raise( - "decimal128 columns require precision and scale", - false, - )); - } - (_, Some(_), _) | (_, _, Some(_)) => { - return Err(raise( - "precision and scale are only valid for decimal128 columns", - false, - )); - } - _ => {} +impl<'py> FromPyObject<'py> for Column { + fn extract_bound(value: &Bound<'py, PyAny>) -> PyResult { + let schema = value.py().import("greptimedb_ingester._types")?; + if !value.is_instance(&schema.getattr("Column")?)? { + return Err(raise("columns must contain Column objects", false)); } + let data_type = match value + .getattr("data_type")? + .getattr("value")? + .extract::()? + { + 0 => ColumnDataType::Boolean, + 1 => ColumnDataType::Int8, + 2 => ColumnDataType::Int16, + 3 => ColumnDataType::Int32, + 4 => ColumnDataType::Int64, + 5 => ColumnDataType::Uint8, + 6 => ColumnDataType::Uint16, + 7 => ColumnDataType::Uint32, + 8 => ColumnDataType::Uint64, + 9 => ColumnDataType::Float32, + 10 => ColumnDataType::Float64, + 11 => ColumnDataType::Binary, + 12 => ColumnDataType::String, + 13 => ColumnDataType::Date, + 15 => ColumnDataType::TimestampSecond, + 16 => ColumnDataType::TimestampMillisecond, + 17 => ColumnDataType::TimestampMicrosecond, + 18 => ColumnDataType::TimestampNanosecond, + 19 => ColumnDataType::TimeSecond, + 20 => ColumnDataType::TimeMillisecond, + 21 => ColumnDataType::TimeMicrosecond, + 22 => ColumnDataType::TimeNanosecond, + 30 => ColumnDataType::Decimal128, + 31 => ColumnDataType::Json, + _ => return Err(raise("Invalid column data type", false)), + }; + let semantic_type = match value + .getattr("semantic_type")? + .getattr("value")? + .extract::()? + { + 0 => SemanticType::Tag, + 1 => SemanticType::Field, + 2 => SemanticType::Timestamp, + _ => return Err(raise("Invalid column semantic type", false)), + }; Ok(Self { - name, + name: value.getattr("name")?.extract()?, data_type, semantic_type, - precision, - scale, + precision: value.getattr("precision")?.extract()?, + scale: value.getattr("scale")?.extract()?, }) } +} - #[getter] - fn name(&self) -> &str { - &self.name - } - - #[getter] - fn data_type(&self) -> ColumnDataType { - self.data_type - } - - #[getter] - fn semantic_type(&self) -> SemanticType { - self.semantic_type - } - - #[getter] - fn precision(&self) -> Option { - self.precision - } - - #[getter] - fn scale(&self) -> Option { - self.scale - } - - fn __repr__(&self) -> String { - match (self.precision, self.scale) { - (Some(precision), Some(scale)) => format!( - "Column({:?}, {}, {}, precision={}, scale={})", - self.name, - self.data_type.__repr__(), - self.semantic_type.__repr__(), - precision, - scale - ), - _ => format!( - "Column({:?}, {}, {})", - self.name, - self.data_type.__repr__(), - self.semantic_type.__repr__() - ), - } +impl Column { + pub fn to_python(&self, py: Python<'_>) -> PyResult> { + let schema = py.import("greptimedb_ingester._types")?; + let dtype = schema + .getattr("ColumnDataType")? + .call1((self.data_type as u8,))?; + let semantic = schema + .getattr("SemanticType")? + .call1((self.semantic_type as u8,))?; + Ok(schema + .getattr("Column")? + .call1((&self.name, dtype, semantic, self.precision, self.scale))? + .unbind()) } } diff --git a/python/greptimedb_ingester/__init__.py b/python/greptimedb_ingester/__init__.py index b0b470a..f9edb76 100644 --- a/python/greptimedb_ingester/__init__.py +++ b/python/greptimedb_ingester/__init__.py @@ -1,40 +1,7 @@ -"""Pure Python facade. Choose optional backends with [grpc] and [postgres].""" - -from importlib.util import find_spec +"""Shared schema for the explicit gRPC and PostgreSQL backends.""" from ._errors import GreptimeError from ._types import Column, ColumnDataType, SemanticType __version__ = "0.1.2" -__all__ = [ - "Client", - "Column", - "ColumnDataType", - "SemanticType", - "GreptimeError", - "__version__", -] - -if find_spec("greptimedb_ingester_grpc") is not None: - # Preserve the existing gRPC Client and native schema classes when installed. - from .grpc import ( - BulkResponse, - BulkWriter, - Client, - Column, - ColumnDataType, - Rows, - SemanticType, - WriteOptions, - ) - - __all__ += ["BulkResponse", "BulkWriter", "Rows", "WriteOptions"] -elif find_spec("psycopg2") is not None: - from .postgres import Client -else: - - class Client: - def __init__(self, *args, **kwargs): - raise ImportError( - "Install greptimedb-ingester[grpc] or greptimedb-ingester[postgres]" - ) +__all__ = ["Column", "ColumnDataType", "SemanticType", "GreptimeError", "__version__"] diff --git a/python/greptimedb_ingester/_errors.py b/python/greptimedb_ingester/_errors.py index 26b8cf4..ead64f2 100644 --- a/python/greptimedb_ingester/_errors.py +++ b/python/greptimedb_ingester/_errors.py @@ -1,12 +1,7 @@ -"""One public error type shared by all installed backends.""" +"""Public error type shared by both backends.""" -try: - from greptimedb_ingester_grpc._native import GreptimeError -except ModuleNotFoundError as exc: - if not exc.name.startswith("greptimedb_ingester_grpc"): - raise - class GreptimeError(Exception): - def __init__(self, message, retriable=False): - super().__init__(message) - self.retriable = retriable +class GreptimeError(Exception): + def __init__(self, message, retriable=False): + super().__init__(message) + self.retriable = retriable diff --git a/python/greptimedb_ingester/_native.py b/python/greptimedb_ingester/_native.py deleted file mode 100644 index 4fa9cf3..0000000 --- a/python/greptimedb_ingester/_native.py +++ /dev/null @@ -1,8 +0,0 @@ -"""Compatibility import for callers of the former native module.""" - -try: - from greptimedb_ingester_grpc._native import * # noqa: F403 -except ModuleNotFoundError as exc: - raise ImportError( - "Install greptimedb-ingester[grpc] for the native extension" - ) from exc diff --git a/python/greptimedb_ingester/_types.py b/python/greptimedb_ingester/_types.py index d2ecc44..67c2e9a 100644 --- a/python/greptimedb_ingester/_types.py +++ b/python/greptimedb_ingester/_types.py @@ -1,4 +1,4 @@ -"""Python schema values matching the gRPC SDK's public names and enum values.""" +"""Public column schema shared by both backends.""" from dataclasses import dataclass from enum import IntEnum @@ -8,9 +8,7 @@ def _schema_error(message): - error = GreptimeError(message) - error.retriable = False - return error + return GreptimeError(message) class _NamedEnum(IntEnum): diff --git a/python/greptimedb_ingester/grpc.py b/python/greptimedb_ingester/grpc.py index e5156db..16f3c7b 100644 --- a/python/greptimedb_ingester/grpc.py +++ b/python/greptimedb_ingester/grpc.py @@ -5,11 +5,7 @@ BulkResponse, BulkWriter, Client, - Column, - ColumnDataType, - GreptimeError, Rows, - SemanticType, WriteOptions, ) except ModuleNotFoundError as exc: @@ -17,10 +13,6 @@ __all__ = [ "Client", - "Column", - "ColumnDataType", - "SemanticType", - "GreptimeError", "WriteOptions", "Rows", "BulkWriter", diff --git a/python/greptimedb_ingester/postgres.py b/python/greptimedb_ingester/postgres.py index dd956c2..f8206d4 100644 --- a/python/greptimedb_ingester/postgres.py +++ b/python/greptimedb_ingester/postgres.py @@ -5,9 +5,7 @@ import math import threading from collections.abc import Mapping -from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from enum import Enum from functools import lru_cache from itertools import islice @@ -20,40 +18,34 @@ from ._errors import GreptimeError +from ._types import Column, ColumnDataType as _DataType, SemanticType as _SemanticType def _error(message, retriable=False, *, acknowledged_rows=None): - error = GreptimeError(message) - error.retriable = retriable + error = GreptimeError(message, retriable=retriable) if acknowledged_rows is not None: error.acknowledged_rows = acknowledged_rows return error -class _SemanticType(Enum): - TAG = "TAG" - FIELD = "FIELD" - TIMESTAMP = "TIMESTAMP" - - -class _DataType(Enum): - BOOLEAN = "BOOLEAN" - INT8 = "TINYINT" - INT16 = "SMALLINT" - INT32 = "INT" - INT64 = "BIGINT" - UINT8 = "TINYINT UNSIGNED" - UINT16 = "SMALLINT UNSIGNED" - UINT32 = "INT UNSIGNED" - UINT64 = "BIGINT UNSIGNED" - FLOAT32 = "FLOAT" - FLOAT64 = "DOUBLE" - STRING = "STRING" - TIMESTAMP_SECOND = "TIMESTAMP(0)" - TIMESTAMP_MILLISECOND = "TIMESTAMP(3)" - TIMESTAMP_MICROSECOND = "TIMESTAMP(6)" - TIMESTAMP_NANOSECOND = "TIMESTAMP(9)" - +_SQL_TYPES = { + _DataType.BOOLEAN: "BOOLEAN", + _DataType.INT8: "TINYINT", + _DataType.INT16: "SMALLINT", + _DataType.INT32: "INT", + _DataType.INT64: "BIGINT", + _DataType.UINT8: "TINYINT UNSIGNED", + _DataType.UINT16: "SMALLINT UNSIGNED", + _DataType.UINT32: "INT UNSIGNED", + _DataType.UINT64: "BIGINT UNSIGNED", + _DataType.FLOAT32: "FLOAT", + _DataType.FLOAT64: "DOUBLE", + _DataType.STRING: "STRING", + _DataType.TIMESTAMP_SECOND: "TIMESTAMP(0)", + _DataType.TIMESTAMP_MILLISECOND: "TIMESTAMP(3)", + _DataType.TIMESTAMP_MICROSECOND: "TIMESTAMP(6)", + _DataType.TIMESTAMP_NANOSECOND: "TIMESTAMP(9)", +} _UNITS = { _DataType.TIMESTAMP_SECOND: 0, @@ -64,60 +56,29 @@ class _DataType(Enum): _EPOCH = datetime(1970, 1, 1, tzinfo=timezone.utc) -@dataclass(frozen=True) -class _Column: - name: str - data_type: _DataType - semantic_type: _SemanticType - - def __post_init__(self): - if not isinstance(self.name, str) or not self.name or "\x00" in self.name: - raise ValueError("Column name must be a nonempty string without NUL") - if not isinstance(self.data_type, _DataType): - raise TypeError("Invalid column data type") - if not isinstance(self.semantic_type, _SemanticType): - raise TypeError("Invalid column semantic type") - if ( - self.semantic_type == _SemanticType.TIMESTAMP - and self.data_type not in _UNITS - ): - raise ValueError("Time index must have a timestamp data type") - if ( - self.semantic_type == _SemanticType.TAG - and self.data_type != _DataType.STRING - ): - raise ValueError("This client supports STRING tags") - - -def _normalize_columns(columns): - result = [] - for column in columns: - try: - name = repr(column.data_type).split(".")[-1] - column_name = column.name - semantic_name = repr(column.semantic_type).split(".")[-1] - except AttributeError as exc: - raise TypeError("columns must contain Column objects") from exc - try: - dtype = _DataType[name] - except KeyError: - raise ValueError(f"PostgreSQL COPY does not support {name}") from None - try: - semantic = _SemanticType[semantic_name] - except KeyError: - raise ValueError(f"Invalid column semantic type: {semantic_name}") from None - result.append(_Column(column_name, dtype, semantic)) - return _columns(result) - - def _columns(columns): columns = tuple(columns) - if not columns or any(not isinstance(c, _Column) for c in columns): + if not columns or any(not isinstance(c, Column) for c in columns): raise ValueError("columns must contain Column objects") if len({c.name for c in columns}) != len(columns): raise ValueError("Duplicate column names") if sum(c.semantic_type == _SemanticType.TIMESTAMP for c in columns) != 1: raise ValueError("Exactly one time index is required") + for column in columns: + if column.data_type not in _SQL_TYPES: + raise ValueError( + f"PostgreSQL COPY does not support {column.data_type.name}" + ) + if ( + column.semantic_type == _SemanticType.TIMESTAMP + and column.data_type not in _UNITS + ): + raise ValueError("Time index must have a timestamp data type") + if ( + column.semantic_type == _SemanticType.TAG + and column.data_type != _DataType.STRING + ): + raise ValueError("This client supports STRING tags") return columns @@ -126,7 +87,7 @@ def _validate_schema(columns, description): for column in columns: found = actual.get(column.name) expected_type = column.data_type.name.replace("_", "") - expected_semantic = column.semantic_type.value + expected_semantic = column.semantic_type.name if found is None: raise _error(f"Table schema mismatch: missing column {column.name!r}") if found[0].upper() != expected_type or found[1] != expected_semantic: @@ -387,7 +348,7 @@ def insert(self, table, columns, rows): neither validation nor transport failures are retried automatically. """ try: - return self._insert(table, _normalize_columns(columns), rows) + return self._insert(table, _columns(columns), rows) except (ValueError, TypeError, OverflowError) as exc: raise _error( f"Input failed after 0 acknowledged rows: {exc}", acknowledged_rows=0 @@ -443,7 +404,8 @@ def _insert(self, table, columns, rows): self._schemas.pop(table, None) definitions = [ sql.SQL("{} {}").format( - sql.Identifier(c.name), sql.SQL(c.data_type.value) + sql.Identifier(c.name), + sql.SQL(_SQL_TYPES[c.data_type]), ) for c in columns ] diff --git a/tests/test_backends.py b/tests/test_backends.py index 85c5714..7b3783a 100644 --- a/tests/test_backends.py +++ b/tests/test_backends.py @@ -34,19 +34,28 @@ def test_public_schema_is_available_without_a_backend(): assert repr(column.semantic_type) == "SemanticType.TAG" -def test_default_client_matches_installed_backend(): - if util.find_spec("greptimedb_ingester_grpc") is not None: - from greptimedb_ingester.grpc import Client, Column, GreptimeError - - assert sdk.Client is Client and sdk.Column is Column - assert sdk.GreptimeError is GreptimeError - elif util.find_spec("psycopg2") is not None: - from greptimedb_ingester.postgres import Client - - assert sdk.Client is Client - else: - with pytest.raises(ImportError, match=r"\[grpc\].*\[postgres\]"): - sdk.Client(["localhost:4003"]) +def test_root_has_no_backend_exports_or_legacy_native_module(): + for name in ["Client", "WriteOptions", "BulkWriter", "Rows", "BulkResponse"]: + assert not hasattr(sdk, name) + assert util.find_spec("greptimedb_ingester._native") is None + + +def test_importing_public_schema_loads_no_backend(): + subprocess.run( + [ + sys.executable, + "-c", + """ +import sys +import greptimedb_ingester +assert 'psycopg2' not in sys.modules +assert 'greptimedb_ingester.postgres' not in sys.modules +assert 'greptimedb_ingester.grpc' not in sys.modules +assert 'greptimedb_ingester_grpc' not in sys.modules +""", + ], + check=True, + ) def test_postgres_errors_are_the_public_error_type(): @@ -79,7 +88,7 @@ def test_grpc_import_does_not_load_the_postgres_driver(): import sys import greptimedb_ingester as sdk from greptimedb_ingester.grpc import Client -assert sdk.Client is Client +assert not hasattr(sdk, "Client") assert "psycopg2" not in sys.modules assert "greptimedb_ingester.postgres" not in sys.modules """ diff --git a/tests/test_convert.py b/tests/test_convert.py index 9f1ffcd..36409cc 100644 --- a/tests/test_convert.py +++ b/tests/test_convert.py @@ -5,15 +5,9 @@ pytest.importorskip("greptimedb_ingester_grpc") -from greptimedb_ingester import ( - Client, - Column, - ColumnDataType, - GreptimeError, - SemanticType, - WriteOptions, -) -from greptimedb_ingester._native import describe_row +from greptimedb_ingester import Column, ColumnDataType, GreptimeError, SemanticType +from greptimedb_ingester.grpc import Client, WriteOptions +from greptimedb_ingester_grpc._native import describe_row from type_samples import ALL_COLUMNS, ALL_ROW, BULK_COLUMNS, BULK_ROW diff --git a/tests/test_integration.py b/tests/test_integration.py index c3e83bc..d9b1723 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -15,7 +15,8 @@ pytest.importorskip("greptimedb_ingester_grpc") -from greptimedb_ingester import Client, Column, ColumnDataType, GreptimeError, SemanticType, WriteOptions +from greptimedb_ingester import Column, ColumnDataType, GreptimeError, SemanticType +from greptimedb_ingester.grpc import Client, WriteOptions pytestmark = pytest.mark.integration @@ -192,6 +193,8 @@ def test_bulk_write_20k(client): COLUMNS, options=WriteOptions(compression="lz4", parallelism=4), ) as writer: + assert writer.columns == COLUMNS + assert all(type(column) is Column for column in writer.columns) rows = writer.alloc_rows(BATCH) assert rows.add_rows(bulk_row(i) for i in range(1, BATCH + 1)) == BATCH response = writer.write(rows) diff --git a/tests/test_postgres.py b/tests/test_postgres.py index fa58431..af8a6e6 100644 --- a/tests/test_postgres.py +++ b/tests/test_postgres.py @@ -274,7 +274,7 @@ def test_all_declared_types_have_an_explicit_supported_or_rejected_case(): "FLOAT64": 1.25, }.get(name, 1) cols = [COLUMNS[1], Column("field", getattr(D, name), S.FIELD)] - normalized = pg._normalize_columns(cols) + normalized = pg._columns(cols) assert pg._row_encoder(normalized)([0, value])[-1] is not None assert pg._row_encoder(normalized)([0, None])[-1] is None @@ -291,7 +291,7 @@ def test_every_integer_width_checks_bounds(name, bits, unsigned): low, high = ( (0, 2**bits - 1) if unsigned else (-(2 ** (bits - 1)), 2 ** (bits - 1) - 1) ) - cols = pg._normalize_columns( + cols = pg._columns( COLUMNS[:2] + [Column("value", getattr(D, name), S.FIELD)] ) encode = pg._row_encoder(cols) @@ -305,7 +305,7 @@ def test_every_integer_width_checks_bounds(name, bits, unsigned): @pytest.mark.parametrize("dtype", [D.FLOAT32, D.FLOAT64]) def test_float_validation(dtype): encode = pg._row_encoder( - pg._normalize_columns([COLUMNS[1], Column("f", dtype, S.FIELD)]) + pg._columns([COLUMNS[1], Column("f", dtype, S.FIELD)]) ) for value in [float("inf"), float("nan"), True, "1"]: with pytest.raises((TypeError, ValueError)): @@ -417,10 +417,6 @@ def test_context_exit_closes_and_preserves_user_exception(database): def test_postgres_only_import_has_no_native_dependency(): - import importlib.util - - if importlib.util.find_spec("greptimedb_ingester_grpc") is not None: - pytest.skip("No-native assertion is exercised by postgres-only CI") code = """ import sys import greptimedb_ingester @@ -681,7 +677,7 @@ def disconnect(*args): ) def test_integer_enum_encodes_as_builtin_int(dtype, enum_type): state = enum_type("State", {"OK": 1}) - columns = pg._normalize_columns([COLUMNS[1], Column("value", dtype, S.FIELD)]) + columns = pg._columns([COLUMNS[1], Column("value", dtype, S.FIELD)]) encoded = pg._row_encoder(columns)([0, state.OK]) assert type(encoded[1]) is int assert encoded[1] == 1 diff --git a/tests/test_postgres_schema_compat.py b/tests/test_postgres_schema_compat.py deleted file mode 100644 index 73f645d..0000000 --- a/tests/test_postgres_schema_compat.py +++ /dev/null @@ -1,41 +0,0 @@ -"""Keep the COPY backend's schema names compatible with gRPC.""" - -import pytest - -pytest.importorskip("greptimedb_ingester_grpc") -from greptimedb_ingester import Column, ColumnDataType, SemanticType -from greptimedb_ingester import _types as pg - - -def test_copy_schema_enum_names_and_values_match_native(): - for native, pure in [ - (ColumnDataType, pg.ColumnDataType), - (SemanticType, pg.SemanticType), - ]: - assert {name for name in dir(native) if name.isupper()} == set(pure.__members__) - for name, member in pure.__members__.items(): - assert getattr(native, name) == member.value - assert repr(getattr(native, name)) == repr(member) - - -def test_copy_column_shape_matches_native(): - native = Column("ts", ColumnDataType.TIMESTAMP_NANOSECOND, SemanticType.TIMESTAMP) - pure = pg.Column( - "ts", pg.ColumnDataType.TIMESTAMP_NANOSECOND, pg.SemanticType.TIMESTAMP - ) - for field in ["name", "precision", "scale"]: - assert getattr(native, field) == getattr(pure, field) - for field in ["data_type", "semantic_type"]: - assert repr(getattr(native, field)) == repr(getattr(pure, field)) - - -def test_installed_copy_client_accepts_native_columns(): - import pytest - - pytest.importorskip("psycopg2") - from greptimedb_ingester.postgres import _normalize_columns, _row_encoder - - native = [Column("ts", ColumnDataType.TIMESTAMP_NANOSECOND, SemanticType.TIMESTAMP)] - assert _row_encoder(_normalize_columns(native))([1]) == [ - "1970-01-01T00:00:00.000000001+00:00" - ] diff --git a/tests/test_shared_schema.py b/tests/test_shared_schema.py new file mode 100644 index 0000000..43c6fd7 --- /dev/null +++ b/tests/test_shared_schema.py @@ -0,0 +1,49 @@ +"""Both backends consume the one public schema.""" + +from types import SimpleNamespace +import pytest +import greptimedb_ingester as sdk + + +def test_postgres_uses_public_columns_directly(): + pytest.importorskip("psycopg2") + from greptimedb_ingester.postgres import _columns, _row_encoder + + column = sdk.Column( + "ts", sdk.ColumnDataType.TIMESTAMP_NANOSECOND, sdk.SemanticType.TIMESTAMP + ) + checked = _columns([column]) + assert checked[0] is column + assert _row_encoder(checked)([1]) == ["1970-01-01T00:00:00.000000001+00:00"] + + +def test_grpc_uses_public_schema_and_error(): + pytest.importorskip("greptimedb_ingester_grpc") + from greptimedb_ingester import grpc + from greptimedb_ingester_grpc import _native + + for name in ["Column", "ColumnDataType", "SemanticType", "GreptimeError"]: + assert name not in grpc.__all__ + assert not hasattr(_native, name) + column = sdk.Column( + "ts", sdk.ColumnDataType.TIMESTAMP_NANOSECOND, sdk.SemanticType.TIMESTAMP + ) + assert _native.describe_row([column], [1]) == ["timestamp_nanosecond:1"] + with pytest.raises(sdk.GreptimeError) as exc: + _native.describe_row([column], ["invalid"]) + assert exc.value.retriable is False + + +def test_grpc_rejects_duck_typed_columns(): + pytest.importorskip("greptimedb_ingester_grpc") + from greptimedb_ingester_grpc._native import describe_row + + column = SimpleNamespace( + name="ts", + data_type=sdk.ColumnDataType.TIMESTAMP_NANOSECOND, + semantic_type=sdk.SemanticType.TIMESTAMP, + precision=None, + scale=None, + ) + with pytest.raises(sdk.GreptimeError, match="Column objects"): + describe_row([column], [1]) diff --git a/tests/test_types.py b/tests/test_types.py index 415de42..9146ef6 100644 --- a/tests/test_types.py +++ b/tests/test_types.py @@ -39,14 +39,14 @@ def test_column_is_immutable_and_decimal_metadata_is_preserved(): column.name = "changed" -def test_compatible_column_protocol_needs_no_native_import(): +def test_postgres_requires_public_column_objects(): pytest.importorskip("psycopg2") - from greptimedb_ingester.postgres import _normalize_columns + from greptimedb_ingester.postgres import _columns column = SimpleNamespace( name="ts", data_type=D.TIMESTAMP_NANOSECOND, semantic_type=S.TIMESTAMP ) - assert _normalize_columns([column])[0].name == "ts" - column.semantic_type = "invalid" - with pytest.raises(ValueError, match="semantic"): - _normalize_columns([column]) + with pytest.raises(ValueError, match="Column objects"): + _columns([column]) + public = Column("ts", D.TIMESTAMP_NANOSECOND, S.TIMESTAMP) + assert _columns([public]) == (public,) From a71be5a44535ad3c8bf88b07290dcdd49fe9fe4f Mon Sep 17 00:00:00 2001 From: paomian Date: Mon, 2 Nov 2026 18:40:12 +0800 Subject: [PATCH 2/2] Refactor and enhance error handling in gRPC and PostgreSQL integration - Update `.gitignore` to ignore all files in the `docs` directory. - Modify README files to reference the new PG getting started guide. - Refactor decimal handling in gRPC to improve precision and scale validation. - Change `exponent` type in `DecimalParts` from `i32` to `i64` for better compatibility. - Enhance error handling in `raise` function to ensure retriable status is correctly set. - Introduce `checked_duration` function to validate duration inputs. - Add comprehensive tests for input validation and error handling in gRPC and PostgreSQL clients. - Ensure schema validation is robust against forged inputs and maintains integrity across operations. - Update PostgreSQL client to handle connection errors gracefully and invalidate cached schemas when necessary. --- .gitignore | 2 +- README.md | 2 +- README.zh-CN.md | 2 +- packages/grpc/src/bulk.rs | 7 +- packages/grpc/src/cell.rs | 13 +- packages/grpc/src/client.rs | 32 +++-- packages/grpc/src/convert.rs | 92 +++++++------- packages/grpc/src/error.rs | 8 +- packages/grpc/src/types.rs | 158 ++++++++++++++++-------- python/greptimedb_ingester/_types.py | 4 +- python/greptimedb_ingester/postgres.py | 50 ++++++-- tests/test_grpc_regressions.py | 160 +++++++++++++++++++++++++ tests/test_postgres.py | 139 +++++++++++++++++++++ tests/test_postgres_integration.py | 16 +++ tests/test_types.py | 10 ++ 15 files changed, 562 insertions(+), 133 deletions(-) create mode 100644 tests/test_grpc_regressions.py diff --git a/.gitignore b/.gitignore index adbe10c..ecbb04b 100644 --- a/.gitignore +++ b/.gitignore @@ -18,4 +18,4 @@ test/ /packages/grpc/target/ .DS_Store -/docs +/docs/* diff --git a/README.md b/README.md index f266c2b..983f06b 100644 --- a/README.md +++ b/README.md @@ -35,7 +35,7 @@ print(affected) `insert` creates missing tables and writes COPY batches (default: 10,000 rows). Batches commit separately. The server must support `COPY FROM STDIN`. -See the [COPY example](examples/postgres_example.py). +See the [PG getting started guide](docs/postgres-getting-started.en.md) and [COPY example](examples/postgres_example.py). ## gRPC diff --git a/README.zh-CN.md b/README.zh-CN.md index 4cc98e2..fc1ce8e 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -35,7 +35,7 @@ print(affected) `insert` 自动创建不存在的表并分批 COPY，默认每批 10,000 行。 各批独立提交。服务端须支持 `COPY FROM STDIN`。 -完整用法见 [COPY 示例](examples/postgres_example.py)。 +完整用法见 [PG 入门指南](docs/postgres-getting-started.md)和 [COPY 示例](examples/postgres_example.py)。 ## gRPC diff --git a/packages/grpc/src/bulk.rs b/packages/grpc/src/bulk.rs index abb2c18..78c19ce 100644 --- a/packages/grpc/src/bulk.rs +++ b/packages/grpc/src/bulk.rs @@ -421,11 +421,8 @@ fn table_schema( crate::types::SemanticType::Field if column.data_type == crate::types::ColumnDataType::Decimal128 => { - schema.add_decimal128_field( - column.name.clone(), - column.precision.expect("precision"), - column.scale.expect("scale"), - ) + let (precision, scale) = column.decimal_params()?; + schema.add_decimal128_field(column.name.clone(), precision, scale) } crate::types::SemanticType::Field => schema.add_field(column.name.clone(), dtype), }; diff --git a/packages/grpc/src/cell.rs b/packages/grpc/src/cell.rs index 2aacfe6..970cb64 100644 --- a/packages/grpc/src/cell.rs +++ b/packages/grpc/src/cell.rs @@ -9,7 +9,7 @@ use crate::error::{raise, RemoteError}; pub struct DecimalParts { pub negative: bool, pub digits: String, - pub exponent: i32, + pub exponent: i64, } #[derive(Clone, Debug)] @@ -102,11 +102,14 @@ pub fn cell_from_py(obj: &Bound<'_, PyAny>) -> PyResult { || obj.is_instance_of::() || obj.is_instance_of::() { + let kwargs = PyDict::new(py); + kwargs.set_item("allow_nan", false)?; let text: String = py .import("json")? .getattr("dumps")? - .call1((obj,))? - .extract()?; + .call((obj,), Some(&kwargs)) + .and_then(|value| value.extract()) + .map_err(|error| raise(format!("invalid JSON: {error}"), false))?; return Ok(Cell::Json(text)); } @@ -131,7 +134,7 @@ fn decimal_cell(obj: &Bound<'_, PyAny>) -> PyResult { if text.is_empty() { text.push('0'); } - let exponent: i32 = parts + let exponent: i64 = parts .getattr("exponent")? .extract() .map_err(|_| raise("Decimal exponent is not an integer", false))?; @@ -145,7 +148,7 @@ fn decimal_cell(obj: &Bound<'_, PyAny>) -> PyResult { fn timestamp_cell(obj: &Bound<'_, PyAny>, datetime_mod: &Bound<'_, PyAny>) -> PyResult { let py = obj.py(); let utc = datetime_mod.getattr("timezone")?.getattr("utc")?; - let aware = if obj.getattr("tzinfo")?.is_none() { + let aware = if obj.getattr("tzinfo")?.is_none() || obj.call_method0("utcoffset")?.is_none() { let kwargs = PyDict::new(py); kwargs.set_item("tzinfo", utc.clone())?; obj.call_method("replace", (), Some(&kwargs))? diff --git a/packages/grpc/src/client.rs b/packages/grpc/src/client.rs index ed02043..5820fd2 100644 --- a/packages/grpc/src/client.rs +++ b/packages/grpc/src/client.rs @@ -12,7 +12,7 @@ use crate::bulk::BulkWriter; use crate::convert::{align, build_delete, build_insert, parse_row}; use crate::error::{raise, RemoteError}; use crate::runtime; -use crate::types::{parse_grpc_compression, Column, WriteOptions}; +use crate::types::{checked_duration, parse_grpc_compression, Column, SemanticType, WriteOptions}; #[pyclass] #[derive(Clone)] @@ -117,7 +117,7 @@ impl Client { if table.is_empty() { return Err(raise("table name must not be empty", false)); } - let aligned = align_rows(&columns, rows)?; + let aligned = align_rows(&columns, rows, "insert")?; let hints = hints.unwrap_or_default(); self.execute(py, move |database| async move { let request = build_insert(&table, &columns, &aligned)?; @@ -147,7 +147,16 @@ impl Client { if table.is_empty() { return Err(raise("table name must not be empty", false)); } - let aligned = align_rows(&columns, rows)?; + if columns + .iter() + .any(|column| column.semantic_type == SemanticType::Field) + { + return Err(raise( + "delete columns must be TAG or TIMESTAMP, not FIELD", + false, + )); + } + let aligned = align_rows(&columns, rows, "delete")?; self.execute(py, move |database| async move { let request = build_delete(&table, &columns, &aligned)?; database.delete(request).await.map_err(RemoteError::from) @@ -212,9 +221,13 @@ impl Client { fn align_rows( columns: &[Column], rows: &Bound<'_, PyAny>, + operation: &str, ) -> PyResult>> { if columns.is_empty() { - return Err(raise("insert requires at least one column", false)); + return Err(raise( + format!("{operation} requires at least one column"), + false, + )); } let mut aligned = Vec::new(); for row in rows.try_iter()? { @@ -222,17 +235,16 @@ fn align_rows( aligned.push(align(columns, &input)?); } if aligned.is_empty() { - return Err(raise("cannot insert an empty row batch", false)); + return Err(raise( + format!("cannot {operation} an empty row batch"), + false, + )); } Ok(aligned) } fn optional_duration(value: Option, name: &str) -> PyResult> { - match value { - None => Ok(None), - Some(secs) if secs >= 0.0 => Ok(Some(Duration::from_secs_f64(secs))), - Some(_) => Err(raise(format!("{name} must be >= 0"), false)), - } + value.map(|secs| checked_duration(secs, name)).transpose() } fn build_ingester_client( diff --git a/packages/grpc/src/convert.rs b/packages/grpc/src/convert.rs index 10776f6..abdaa88 100644 --- a/packages/grpc/src/convert.rs +++ b/packages/grpc/src/convert.rs @@ -105,8 +105,9 @@ pub fn column_schema(column: &Column) -> Result { SemanticType::Field => field(&column.name, dtype), }; if column.data_type == ColumnDataType::Decimal128 { - let precision = i32::from(column.precision.expect("decimal precision")); - let scale = i32::from(column.scale.expect("decimal scale")); + let (precision, scale) = column.decimal_params()?; + let precision = i32::from(precision); + let scale = i32::from(scale); schema.datatype_extension = Some(ColumnDataTypeExtension { type_ext: Some(TypeExt::DecimalType(DecimalTypeExtension { precision, @@ -293,9 +294,7 @@ pub fn to_proto( ColumnDataType::TimeNanosecond => { time_nanosecond_value(expect_int(cell, "time_nanosecond")?) } - ColumnDataType::Decimal128 => { - decimal128_value(expect_decimal(cell, column.scale.expect("scale"))?) - } + ColumnDataType::Decimal128 => decimal128_value(expect_decimal(cell, column)?), ColumnDataType::Json => { let text = expect_json_text(cell)?; serde_json::from_str::(&text) @@ -351,9 +350,7 @@ pub fn to_bulk(column: &Column, cell: &Cell) -> Result { ColumnDataType::TimeNanosecond => { Value::TimeNanosecond(expect_int(cell, "time_nanosecond")?) } - ColumnDataType::Decimal128 => { - Value::Decimal128(expect_decimal(cell, column.scale.expect("scale"))?) - } + ColumnDataType::Decimal128 => Value::Decimal128(expect_decimal(cell, column)?), ColumnDataType::Json => Value::Binary(expect_bytes(cell)?), }; Ok(value) @@ -521,49 +518,60 @@ fn expect_epoch(cell: &Cell, unit_nanos: i64) -> Result { } fn scale_timestamp(secs: i64, subsec_nanos: u32, unit_nanos: i64) -> Result { - let total = secs - .checked_mul(1_000_000_000) - .and_then(|value| value.checked_add(i64::from(subsec_nanos))) - .ok_or_else(|| RemoteError::new("timestamp value is out of range", false))?; - total - .checked_div(unit_nanos) - .ok_or_else(|| RemoteError::new("timestamp value is out of range", false)) + let total = i128::from(secs) * 1_000_000_000 + i128::from(subsec_nanos); + fit_int(total.div_euclid(i128::from(unit_nanos)), "timestamp") } -fn expect_decimal(cell: &Cell, scale: i8) -> Result { - match cell { - Cell::Int(value) => Ok(*value), - Cell::Decimal(parts) => scale_decimal(parts, scale), - other => Err(type_mismatch("decimal128", other)), +fn expect_decimal(cell: &Cell, column: &Column) -> Result { + let (precision, scale) = column.decimal_params()?; + let value = match cell { + Cell::Int(value) => *value, + Cell::Decimal(parts) => scale_decimal(parts, scale)?, + other => return Err(type_mismatch("decimal128", other)), + }; + let limit = 10i128.pow(u32::from(precision)); + if value <= -limit || value >= limit { + return Err(RemoteError::new( + format!("decimal128 value exceeds precision {precision}"), + false, + )); } + Ok(value) } fn scale_decimal(parts: &DecimalParts, scale: i8) -> Result { - let coefficient: i128 = parts - .digits - .parse() - .map_err(|_| RemoteError::new("decimal128 value is out of range", false))?; - let shift = parts.exponent + i32::from(scale); - let scaled = if shift >= 0 { - let factor = 10i128 - .checked_pow(shift as u32) - .ok_or_else(|| RemoteError::new("decimal128 value is out of range", false))?; - coefficient - .checked_mul(factor) - .ok_or_else(|| RemoteError::new("decimal128 value is out of range", false))? + let digits = parts.digits.trim_start_matches('0'); + if digits.is_empty() { + return Ok(0); + } + let shift = i128::from(parts.exponent) + i128::from(scale); + let text = if shift < 0 { + let removed = usize::try_from(-shift).unwrap_or(usize::MAX); + if removed >= digits.len() || !digits[digits.len() - removed..].bytes().all(|b| b == b'0') { + return Err(RemoteError::new( + format!("decimal has more fractional digits than scale {scale}"), + false, + )); + } + digits[..digits.len() - removed].to_owned() } else { - return Err(RemoteError::new( - format!("decimal has more fractional digits than scale {scale}"), - false, - )); + // Decimal128 has at most 38 significant decimal digits. Bound the work + // before allocating, even for Decimal exponents near Python's limit. + if shift > 38 || digits.len() > 38 - shift as usize { + return Err(RemoteError::new("decimal128 value is out of range", false)); + } + let mut text = digits.to_owned(); + text.extend(std::iter::repeat('0').take(shift as usize)); + text }; - if parts.negative { - scaled - .checked_neg() - .ok_or_else(|| RemoteError::new("decimal128 value is out of range", false)) + let coefficient: i128 = text + .parse() + .map_err(|_| RemoteError::new("decimal128 value is out of range", false))?; + Ok(if parts.negative { + -coefficient } else { - Ok(scaled) - } + coefficient + }) } #[pyfunction] diff --git a/packages/grpc/src/error.rs b/packages/grpc/src/error.rs index b73911f..996023e 100644 --- a/packages/grpc/src/error.rs +++ b/packages/grpc/src/error.rs @@ -31,10 +31,14 @@ pub fn raise(message: impl Into, retriable: bool) -> PyErr { match py .import("greptimedb_ingester._errors") .and_then(|module| module.getattr("GreptimeError")) - .and_then(|class| class.call1((message, retriable))) + .and_then(|class| class.call1((&message, retriable))) { Ok(instance) => PyErr::from_value(instance), - Err(error) => error, + Err(_) => { + let error = pyo3::exceptions::PyRuntimeError::new_err(message); + let _ = error.value(py).setattr("retriable", retriable); + error + } } }) } diff --git a/packages/grpc/src/types.rs b/packages/grpc/src/types.rs index d6c2d8c..e574f79 100644 --- a/packages/grpc/src/types.rs +++ b/packages/grpc/src/types.rs @@ -1,6 +1,7 @@ use std::time::Duration; use pyo3::prelude::*; +use pyo3::types::PyInt; use crate::error::raise; use greptimedb_ingester::api::v1::ColumnDataType as ProtoType; @@ -87,58 +88,104 @@ impl<'py> FromPyObject<'py> for Column { if !value.is_instance(&schema.getattr("Column")?)? { return Err(raise("columns must contain Column objects", false)); } - let data_type = match value - .getattr("data_type")? - .getattr("value")? - .extract::()? - { - 0 => ColumnDataType::Boolean, - 1 => ColumnDataType::Int8, - 2 => ColumnDataType::Int16, - 3 => ColumnDataType::Int32, - 4 => ColumnDataType::Int64, - 5 => ColumnDataType::Uint8, - 6 => ColumnDataType::Uint16, - 7 => ColumnDataType::Uint32, - 8 => ColumnDataType::Uint64, - 9 => ColumnDataType::Float32, - 10 => ColumnDataType::Float64, - 11 => ColumnDataType::Binary, - 12 => ColumnDataType::String, - 13 => ColumnDataType::Date, - 15 => ColumnDataType::TimestampSecond, - 16 => ColumnDataType::TimestampMillisecond, - 17 => ColumnDataType::TimestampMicrosecond, - 18 => ColumnDataType::TimestampNanosecond, - 19 => ColumnDataType::TimeSecond, - 20 => ColumnDataType::TimeMillisecond, - 21 => ColumnDataType::TimeMicrosecond, - 22 => ColumnDataType::TimeNanosecond, - 30 => ColumnDataType::Decimal128, - 31 => ColumnDataType::Json, - _ => return Err(raise("Invalid column data type", false)), - }; - let semantic_type = match value - .getattr("semantic_type")? - .getattr("value")? - .extract::()? - { - 0 => SemanticType::Tag, - 1 => SemanticType::Field, - 2 => SemanticType::Timestamp, - _ => return Err(raise("Invalid column semantic type", false)), - }; - Ok(Self { - name: value.getattr("name")?.extract()?, - data_type, - semantic_type, - precision: value.getattr("precision")?.extract()?, - scale: value.getattr("scale")?.extract()?, - }) + let result = (|| { + let data_type_obj = value.getattr("data_type")?; + if !data_type_obj.is_instance(&schema.getattr("ColumnDataType")?)? { + return Err(raise("data_type must be ColumnDataType", false)); + } + let semantic_obj = value.getattr("semantic_type")?; + if !semantic_obj.is_instance(&schema.getattr("SemanticType")?)? { + return Err(raise("semantic_type must be SemanticType", false)); + } + for name in ["precision", "scale"] { + let field = value.getattr(name)?; + if !field.is_none() && !field.is_exact_instance_of::() { + return Err(raise( + format!("decimal128 {name} must be an integer"), + false, + )); + } + } + let data_type = match data_type_obj.getattr("value")?.extract::()? { + 0 => ColumnDataType::Boolean, + 1 => ColumnDataType::Int8, + 2 => ColumnDataType::Int16, + 3 => ColumnDataType::Int32, + 4 => ColumnDataType::Int64, + 5 => ColumnDataType::Uint8, + 6 => ColumnDataType::Uint16, + 7 => ColumnDataType::Uint32, + 8 => ColumnDataType::Uint64, + 9 => ColumnDataType::Float32, + 10 => ColumnDataType::Float64, + 11 => ColumnDataType::Binary, + 12 => ColumnDataType::String, + 13 => ColumnDataType::Date, + 15 => ColumnDataType::TimestampSecond, + 16 => ColumnDataType::TimestampMillisecond, + 17 => ColumnDataType::TimestampMicrosecond, + 18 => ColumnDataType::TimestampNanosecond, + 19 => ColumnDataType::TimeSecond, + 20 => ColumnDataType::TimeMillisecond, + 21 => ColumnDataType::TimeMicrosecond, + 22 => ColumnDataType::TimeNanosecond, + 30 => ColumnDataType::Decimal128, + 31 => ColumnDataType::Json, + _ => return Err(raise("Invalid column data type", false)), + }; + let semantic_type = match semantic_obj.getattr("value")?.extract::()? { + 0 => SemanticType::Tag, + 1 => SemanticType::Field, + 2 => SemanticType::Timestamp, + _ => return Err(raise("Invalid column semantic type", false)), + }; + let column = Self { + name: value.getattr("name")?.extract()?, + data_type, + semantic_type, + precision: value.getattr("precision")?.extract()?, + scale: value.getattr("scale")?.extract()?, + }; + column.validate().map_err(PyErr::from)?; + Ok(column) + })(); + result.map_err(|error: PyErr| raise(format!("Invalid column schema: {error}"), false)) } } impl Column { + fn validate(&self) -> Result<(), crate::error::RemoteError> { + use crate::error::RemoteError; + if self.name.is_empty() || self.name.contains('\0') { + return Err(RemoteError::new( + "column name must be nonempty without NUL", + false, + )); + } + if self.data_type == ColumnDataType::Decimal128 { + self.decimal_params()?; + } else if self.precision.is_some() || self.scale.is_some() { + return Err(RemoteError::new( + "precision and scale are only valid for decimal128 columns", + false, + )); + } + Ok(()) + } + + pub fn decimal_params(&self) -> Result<(u8, i8), crate::error::RemoteError> { + use crate::error::RemoteError; + match (self.precision, self.scale) { + (Some(precision @ 1..=38), Some(scale)) if scale >= 0 && scale as u8 <= precision => { + Ok((precision, scale)) + } + _ => Err(RemoteError::new( + "decimal128 requires precision between 1 and 38 and scale between 0 and precision", + false, + )), + } + } + pub fn to_python(&self, py: Python<'_>) -> PyResult> { let schema = py.import("greptimedb_ingester._types")?; let dtype = schema @@ -196,6 +243,15 @@ pub fn parse_grpc_compression(value: &str) -> PyResult { } } +pub fn checked_duration(secs: f64, name: &str) -> PyResult { + Duration::try_from_secs_f64(secs).map_err(|_| { + raise( + format!("{name} must be finite, >= 0, and within Duration range"), + false, + ) + }) +} + #[pyclass(name = "WriteOptions")] #[derive(Clone, Debug)] pub struct WriteOptions { @@ -209,15 +265,13 @@ impl WriteOptions { #[new] #[pyo3(signature = (compression = "lz4", timeout_secs = 60.0, parallelism = 4))] fn new(compression: &str, timeout_secs: f64, parallelism: usize) -> PyResult { - if timeout_secs < 0.0 { - return Err(raise("timeout_secs must be >= 0", false)); - } + let timeout = checked_duration(timeout_secs, "timeout_secs")?; if parallelism == 0 { return Err(raise("parallelism must be >= 1", false)); } Ok(Self { compression: BulkCompression::parse(compression)?, - timeout: Duration::from_secs_f64(timeout_secs), + timeout, parallelism, }) } diff --git a/python/greptimedb_ingester/_types.py b/python/greptimedb_ingester/_types.py index 67c2e9a..7b6fc22 100644 --- a/python/greptimedb_ingester/_types.py +++ b/python/greptimedb_ingester/_types.py @@ -1,7 +1,7 @@ """Public column schema shared by both backends.""" from dataclasses import dataclass -from enum import IntEnum +from enum import Enum from ._errors import GreptimeError @@ -11,7 +11,7 @@ def _schema_error(message): return GreptimeError(message) -class _NamedEnum(IntEnum): +class _NamedEnum(Enum): def __repr__(self): return f"{type(self).__name__}.{self.name}" diff --git a/python/greptimedb_ingester/postgres.py b/python/greptimedb_ingester/postgres.py index f8206d4..a47760a 100644 --- a/python/greptimedb_ingester/postgres.py +++ b/python/greptimedb_ingester/postgres.py @@ -65,6 +65,7 @@ def _columns(columns): if sum(c.semantic_type == _SemanticType.TIMESTAMP for c in columns) != 1: raise ValueError("Exactly one time index is required") for column in columns: + Column.__post_init__(column) if column.data_type not in _SQL_TYPES: raise ValueError( f"PostgreSQL COPY does not support {column.data_type.name}" @@ -134,6 +135,8 @@ def encode(value): raise TypeError( "Timestamp must be an integer in the column's unit or an aware datetime" ) + if not -(2**63) <= units < 2**63: + raise ValueError("Timestamp value is outside INT64 range") seconds, fraction = divmod(units, unit) text = _second_text(seconds) return ( @@ -327,6 +330,12 @@ def _connection(self): self._conn.autocommit = True return self._conn + def _discard_connection(self): + connection, self._conn = self._conn, None + self._schemas.clear() + if connection is not None: + connection.close() + def health_check(self): with self._lock: try: @@ -334,6 +343,8 @@ def health_check(self): cur.execute("SELECT 1") cur.fetchone() except psycopg2.Error as exc: + if isinstance(exc, (psycopg2.OperationalError, psycopg2.InterfaceError)): + self._discard_connection() raise _error( str(exc), isinstance( @@ -347,26 +358,30 @@ def insert(self, table, columns, rows): Batches are independently acknowledged. Failure may leave partial writes; neither validation nor transport failures are retried automatically. """ - try: - return self._insert(table, _columns(columns), rows) - except (ValueError, TypeError, OverflowError) as exc: - raise _error( - f"Input failed after 0 acknowledged rows: {exc}", acknowledged_rows=0 - ) from exc + with self._lock: + if self._closed: + raise _error("Client is closed", acknowledged_rows=0) + try: + columns = _columns(columns) + except Exception as exc: + raise _error( + f"Input failed after 0 acknowledged rows: {exc}", acknowledged_rows=0 + ) from exc + return self._insert(table, columns, rows) def _insert(self, table, columns, rows): - encode_row = _row_encoder(columns) - identifier = _table(table) - command = sql.SQL("COPY {} ({}) FROM STDIN WITH (FORMAT csv)").format( - identifier, sql.SQL(", ").join(sql.Identifier(c.name) for c in columns) - ) - copy_command = None try: + encode_row = _row_encoder(columns) + identifier = _table(table) + command = sql.SQL("COPY {} ({}) FROM STDIN WITH (FORMAT csv)").format( + identifier, sql.SQL(", ").join(sql.Identifier(c.name) for c in columns) + ) iterator = iter(rows) except Exception as exc: raise _error( f"Input failed after 0 acknowledged rows: {exc}", acknowledged_rows=0 ) from exc + copy_command = None affected = 0 with self._lock: if self._closed: @@ -448,7 +463,18 @@ def _insert(self, table, columns, rows): f"Expected {len(batch)} acknowledged rows, received {cur.rowcount}", acknowledged_rows=affected, ) + except GreptimeError as exc: + if not hasattr(exc, "acknowledged_rows"): + raise _error( + f"COPY failed after {affected} acknowledged rows: {exc}", + exc.retriable, + acknowledged_rows=affected, + ) from exc + raise except psycopg2.Error as exc: + self._schemas.pop(table, None) + if isinstance(exc, (psycopg2.OperationalError, psycopg2.InterfaceError)): + self._discard_connection() raise _error( f"COPY failed after {affected} acknowledged rows; current batch may be partial: {exc}", isinstance( diff --git a/tests/test_grpc_regressions.py b/tests/test_grpc_regressions.py new file mode 100644 index 0000000..1561b7a --- /dev/null +++ b/tests/test_grpc_regressions.py @@ -0,0 +1,160 @@ +"""Input validation and conversion regressions, without a database.""" + +from datetime import datetime, timezone, tzinfo +from decimal import Decimal + +import pytest + +pytest.importorskip("greptimedb_ingester_grpc") +from greptimedb_ingester import Column, ColumnDataType as D, GreptimeError, SemanticType as S +from greptimedb_ingester.grpc import Client, WriteOptions +from greptimedb_ingester_grpc._native import describe_row + + +@pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf"), 1e30, -1.0]) +@pytest.mark.parametrize("option", ["timeout_secs", "connect_timeout_secs", "write_timeout"]) +def test_invalid_timeout_is_a_nonretriable_error(value, option): + with pytest.raises(GreptimeError, match="timeout") as exc: + if option == "write_timeout": + WriteOptions(timeout_secs=value) + else: + Client(["127.0.0.1:4001"], **{option: value}) + assert exc.value.retriable is False + + +@pytest.mark.parametrize("changes", [ + {"precision": None}, {"scale": None}, {"precision": 256}, + {"precision": True}, {"scale": -1}, {"scale": 11}, + {"name": ""}, {"name": "bad\x00name"}, + {"data_type": S.FIELD}, {"semantic_type": D.INT8}, +]) +@pytest.mark.parametrize("bulk", [False, True]) +def test_forged_schema_is_rejected_without_panicking(changes, bulk): + column = Column("price", D.DECIMAL128, S.FIELD, precision=10, scale=2) + for name, value in changes.items(): + object.__setattr__(column, name, value) + with pytest.raises(GreptimeError) as exc: + describe_row([column], [Decimal("1.00")], bulk=bulk) + assert exc.value.retriable is False + + +@pytest.mark.parametrize("operation", ["insert", "delete", "bulk_writer"]) +def test_forged_schema_is_rejected_at_every_client_entry(operation): + column = Column("price", D.DECIMAL128, S.FIELD, precision=10, scale=2) + object.__setattr__(column, "scale", None) + client = Client(["127.0.0.1:1"], timeout_secs=0.01, connect_timeout_secs=0.01) + with pytest.raises(GreptimeError, match="decimal") as exc: + if operation == "bulk_writer": + client.bulk_writer("t", [column]) + else: + getattr(client, operation)("t", [column], [[1]]) + assert exc.value.retriable is False + + +@pytest.mark.parametrize("bulk", [False, True]) +@pytest.mark.parametrize("dtype,expected", [ + (D.TIMESTAMP_SECOND, -1), (D.TIMESTAMP_MILLISECOND, -1), + (D.TIMESTAMP_MICROSECOND, -1), (D.TIMESTAMP_NANOSECOND, -1000), +]) +def test_pre_epoch_datetime_rounds_down(dtype, expected, bulk): + column = Column("ts", dtype, S.TIMESTAMP) + value = datetime(1969, 12, 31, 23, 59, 59, 999999, tzinfo=timezone.utc) + assert describe_row([column], [value], bulk=bulk) == [f"{dtype.name.lower()}:{expected}"] + + +@pytest.mark.parametrize("bulk", [False, True]) +@pytest.mark.parametrize("year", [1, 2500, 9999]) +@pytest.mark.parametrize("dtype,multiplier", [ + (D.TIMESTAMP_SECOND, 1), (D.TIMESTAMP_MILLISECOND, 1000), + (D.TIMESTAMP_MICROSECOND, 1000000), +]) +def test_datetime_range_depends_on_column_unit(year, dtype, multiplier, bulk): + value = datetime(year, 1, 1, tzinfo=timezone.utc) + delta = value - datetime(1970, 1, 1, tzinfo=timezone.utc) + expected = delta.days * 86400 * multiplier + column = Column("ts", dtype, S.TIMESTAMP) + assert describe_row([column], [value], bulk=bulk) == [f"{dtype.name.lower()}:{expected}"] + + +@pytest.mark.parametrize("bulk", [False, True]) +def test_datetime_nanoseconds_still_reject_overflow(bulk): + column = Column("ts", D.TIMESTAMP_NANOSECOND, S.TIMESTAMP) + with pytest.raises(GreptimeError, match="out of range"): + describe_row([column], [datetime(2500, 1, 1, tzinfo=timezone.utc)], bulk=bulk) + + +class NaiveTimezone(tzinfo): + def utcoffset(self, dt): + return None + + +@pytest.mark.parametrize("bulk", [False, True]) +def test_datetime_with_no_offset_follows_naive_utc_rule(bulk): + column = Column("ts", D.TIMESTAMP_SECOND, S.TIMESTAMP) + value = datetime(1970, 1, 1, tzinfo=NaiveTimezone()) + assert describe_row([column], [value], bulk=bulk) == ["timestamp_second:0"] + + +@pytest.mark.parametrize("bulk", [False, True]) +@pytest.mark.parametrize("value,expected", [ + (Decimal("12.340"), 1234), (Decimal("-12.34000"), -1234), + (Decimal("0E+1000000"), 0), (Decimal("0E-1000000"), 0), + (Decimal("0E+9999999999"), 0), (Decimal("0E-9999999999"), 0), + (Decimal("1." + "0" * 100), 100), + (Decimal("99.99"), 9999), (Decimal("-99.99"), -9999), + (9999, 9999), +]) +def test_decimal_exact_scale_and_zero(value, expected, bulk): + column = Column("price", D.DECIMAL128, S.FIELD, precision=4, scale=2) + assert describe_row([column], [value], bulk=bulk) == [f"decimal128:{expected}"] + + +@pytest.mark.parametrize("bulk", [False, True]) +@pytest.mark.parametrize("value", [Decimal("100.00"), Decimal("-100.00"), 10000, -10000, 2**127 - 1]) +def test_decimal_enforces_declared_precision(value, bulk): + column = Column("price", D.DECIMAL128, S.FIELD, precision=4, scale=2) + with pytest.raises(GreptimeError, match="precision") as exc: + describe_row([column], [value], bulk=bulk) + assert exc.value.retriable is False + + +@pytest.mark.parametrize("value", [{"x": object()}, {"x": float("nan")}, {"x": Decimal("1")}]) +def test_json_serialization_errors_use_public_error(value): + column = Column("json", D.JSON, S.FIELD) + with pytest.raises(GreptimeError, match="JSON") as exc: + describe_row([column], [value]) + assert exc.value.retriable is False + + +def test_circular_json_uses_public_error(): + value = [] + value.append(value) + with pytest.raises(GreptimeError, match="JSON"): + describe_row([Column("json", D.JSON, S.FIELD)], [value]) + + +@pytest.mark.parametrize("columns,rows,message", [ + ([], [[1]], "delete requires at least one column"), + ([Column("ts", D.TIMESTAMP_SECOND, S.TIMESTAMP)], [], "cannot delete an empty row batch"), + ([Column("v", D.INT64, S.FIELD)], [[1]], "FIELD"), +]) +def test_delete_rejects_invalid_input_before_connecting(columns, rows, message): + client = Client(["127.0.0.1:1"], timeout_secs=0.01, connect_timeout_secs=0.01) + with pytest.raises(GreptimeError, match=message) as exc: + client.delete("t", columns, rows) + assert exc.value.retriable is False + + +@pytest.mark.parametrize("value", [D.INT8, S.TAG]) +def test_schema_enums_are_not_integer_cells(value): + with pytest.raises(GreptimeError, match="unsupported Python value"): + describe_row([Column("n", D.INT64, S.FIELD)], [value]) + + +def test_error_fallback_retains_original_message_and_retriable(monkeypatch): + import sys + + monkeypatch.setitem(sys.modules, "greptimedb_ingester._errors", None) + with pytest.raises(RuntimeError, match="timeout_secs") as exc: + WriteOptions(timeout_secs=float("nan")) + assert exc.value.retriable is False diff --git a/tests/test_postgres.py b/tests/test_postgres.py index af8a6e6..cf706db 100644 --- a/tests/test_postgres.py +++ b/tests/test_postgres.py @@ -865,3 +865,142 @@ def rows(): assert consumed == list(range(failed_batch * 2)) cur.copy_expert.side_effect = successful_copy assert client.insert("ack_mismatch", COLUMNS, [["a", 10, 10]]) == 1 + + +@pytest.mark.parametrize("invalid_schema", [False, True]) +def test_closed_client_does_not_touch_inputs(database, invalid_schema): + class Untouched: + def __iter__(self): + pytest.fail("closed client must not iterate input") + + client = pg.Client("localhost:4003") + client.close() + with pytest.raises(GreptimeError, match="Client is closed") as exc: + client.insert("t", Untouched() if invalid_schema else COLUMNS, Untouched()) + assert exc.value.acknowledged_rows == 0 + assert exc.value.retriable is False + database[2].assert_not_called() + + +def test_schema_iterator_failure_is_a_public_error(database): + def columns(): + yield COLUMNS[0] + raise LookupError("schema source failed") + + with pytest.raises(GreptimeError, match="schema source failed") as exc: + pg.Client("localhost:4003").insert("t", columns(), []) + assert exc.value.acknowledged_rows == 0 + assert exc.value.retriable is False + assert isinstance(exc.value.__cause__, LookupError) + database[2].assert_not_called() + + +@pytest.mark.parametrize("value", [2**63, -(2**63) - 1, datetime(2500, 1, 1, tzinfo=timezone.utc)]) +def test_timestamp_overflow_is_rejected_before_copy(database, value): + with pytest.raises(GreptimeError, match="range") as exc: + pg.Client("localhost:4003").insert("t", COLUMNS, [["a", value, 1]]) + assert exc.value.acknowledged_rows == 0 + database[2].assert_not_called() + + +def test_copy_sql_failure_invalidates_cached_schema_without_retry(database): + _, cur, connect, payloads = database + with pg.Client("localhost:4003") as client: + assert client.insert("recreated", COLUMNS, [["a", 0, 1]]) == 1 + success = cur.copy_expert.side_effect + cur.copy_expert.side_effect = pg.psycopg2.ProgrammingError("table missing") + with pytest.raises(GreptimeError) as exc: + client.insert("recreated", COLUMNS, [["a", 1, 2]]) + assert exc.value.acknowledged_rows == 0 + assert "recreated" not in client._schemas + assert cur.copy_expert.call_count == 2 + cur.copy_expert.side_effect = success + assert client.insert("recreated", COLUMNS, [["a", 2, 3]]) == 1 + assert cur.fetchall.call_count == 2 + assert connect.call_count == 1 + assert len(payloads) == 2 + + +@pytest.mark.parametrize("operation", ["insert", "health_check"]) +@pytest.mark.parametrize("error_type", [pg.psycopg2.OperationalError, pg.psycopg2.InterfaceError]) +def test_transport_failure_discards_even_a_nominally_open_connection(database, operation, error_type): + conn, cur, connect, _ = database + healthy = MagicMock(closed=0) + healthy_cur = healthy.cursor.return_value.__enter__.return_value + healthy_cur.rowcount = 1 + healthy_cur.fetchall.return_value = TABLE_SCHEMA + connect.side_effect = [conn, healthy] + with pg.Client("localhost:4003") as client: + assert client.insert("t", COLUMNS, [["a", 0, 1]]) == 1 + if operation == "insert": + cur.copy_expert.side_effect = error_type("connection unusable") + call = lambda: client.insert("t", COLUMNS, [["a", 1, 2]]) + else: + cur.execute.side_effect = error_type("connection unusable") + call = client.health_check + with pytest.raises(GreptimeError) as exc: + call() + assert exc.value.retriable is True + assert client._conn is None + assert client._schemas == {} + assert conn.closed == 0 # driver has not marked the socket closed + assert connect.call_count == 1 # never retry the failed operation + assert client.insert("t", COLUMNS, [["a", 2, 3]]) == 1 + assert connect.call_count == 2 + conn.close.assert_called_once() + healthy_cur.fetchall.assert_called_once() + + +@pytest.mark.parametrize("value", [D.INT8, S.TAG]) +def test_schema_enums_cannot_be_written_as_numbers(database, value): + with pytest.raises(GreptimeError): + pg.Client("localhost:4003").insert("t", COLUMNS, [["a", 0, value]]) + database[2].assert_not_called() + + + +def test_schema_validation_failure_reports_acknowledged_rows(database): + _, cur, _, _ = database + cur.fetchall.return_value = TABLE_SCHEMA[1:] + with pytest.raises(GreptimeError, match="after 0 acknowledged rows") as exc: + pg.Client("localhost:4003").insert("t", COLUMNS, [["a", 0, 1]]) + assert exc.value.acknowledged_rows == 0 + assert exc.value.retriable is False + cur.copy_expert.assert_not_called() + + +@pytest.mark.parametrize("name,value", [ + ("name", ""), ("name", "bad\x00name"), ("data_type", S.FIELD), + ("semantic_type", D.BOOLEAN), ("precision", 1), +]) +def test_forged_public_column_is_revalidated_before_copy(database, name, value): + tag = Column("host", D.STRING, S.TAG) + object.__setattr__(tag, name, value) + with pytest.raises(GreptimeError) as exc: + pg.Client("localhost:4003").insert("t", [tag, *COLUMNS[1:]], [["a", 0, 1]]) + assert exc.value.acknowledged_rows == 0 + assert exc.value.retriable is False + database[2].assert_not_called() + + +def test_schema_failure_after_reconnect_preserves_acknowledged_prefix(database): + conn, cur, connect, _ = database + healthy = MagicMock(closed=0) + healthy_cur = healthy.cursor.return_value.__enter__.return_value + healthy_cur.fetchall.return_value = TABLE_SCHEMA[1:] + connect.side_effect = [conn, healthy] + successful_copy = cur.copy_expert.side_effect + + def copy_then_disconnect(command, data): + successful_copy(command, data) + conn.closed = 1 + + cur.copy_expert.side_effect = copy_then_disconnect + with pg.Client("localhost:4003", batch_size=1) as client: + with pytest.raises(GreptimeError, match="after 1 acknowledged rows") as exc: + client.insert("t", COLUMNS, [["a", 0, 1], ["b", 1, 2]]) + assert exc.value.acknowledged_rows == 1 + assert exc.value.retriable is False + assert "schema mismatch" in str(exc.value) + assert "t" not in client._schemas + healthy_cur.copy_expert.assert_not_called() diff --git a/tests/test_postgres_integration.py b/tests/test_postgres_integration.py index 2360267..cd210d8 100644 --- a/tests/test_postgres_integration.py +++ b/tests/test_postgres_integration.py @@ -487,3 +487,19 @@ def rows(): sql.SQL("SELECT seq FROM {} ORDER BY ts").format(sql.Identifier(table)), ) == [(0,), (1,)] assert client.insert(table, COLUMNS[:3], [["a", BASE_MS + 3, 3]]) == 1 + + +def test_external_drop_invalidates_cache_and_next_call_recreates(client, query, table): + columns = COLUMNS[:3] + assert client.insert(table, columns, [["a", BASE_MS, 1]]) == 1 + with query.cursor() as cur: + cur.execute(sql.SQL("DROP TABLE {}").format(sql.Identifier(table))) + with pytest.raises(GreptimeError, match="COPY failed") as exc: + client.insert(table, columns, [["b", BASE_MS + 1, 2]]) + assert exc.value.acknowledged_rows == 0 + assert table not in client._schemas + # Recreate only on the caller's next operation; never replay the failed row. + assert client.insert(table, columns, [["c", BASE_MS + 2, 3]]) == 1 + assert records( + query, sql.SQL("SELECT host, seq FROM {}").format(sql.Identifier(table)) + ) == [("c", 3)] diff --git a/tests/test_types.py b/tests/test_types.py index 9146ef6..c67f41a 100644 --- a/tests/test_types.py +++ b/tests/test_types.py @@ -50,3 +50,13 @@ def test_postgres_requires_public_column_objects(): _columns([column]) public = Column("ts", D.TIMESTAMP_NANOSECOND, S.TIMESTAMP) assert _columns([public]) == (public,) + + +def test_schema_enums_have_names_and_are_not_integer_values(): + assert str(D.STRING) == "ColumnDataType.STRING" + assert f"{D.STRING}" == "ColumnDataType.STRING" + assert str(S.TAG) == "SemanticType.TAG" + assert f"{S.TAG}" == "SemanticType.TAG" + assert len({D.BOOLEAN, S.TAG, 0, False}) == 3 + assert not isinstance(D.STRING, int) + assert D.STRING.value == 12