ob-postgres 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- ob_postgres-0.1.0/.gitignore +35 -0
- ob_postgres-0.1.0/CLAUDE.md +109 -0
- ob_postgres-0.1.0/LICENSE +70 -0
- ob_postgres-0.1.0/PKG-INFO +37 -0
- ob_postgres-0.1.0/README.md +15 -0
- ob_postgres-0.1.0/pyproject.toml +52 -0
- ob_postgres-0.1.0/src/ob_postgres/__init__.py +157 -0
- ob_postgres-0.1.0/src/ob_postgres/compiler.py +8 -0
- ob_postgres-0.1.0/src/ob_postgres/connection.py +65 -0
- ob_postgres-0.1.0/src/ob_postgres/cursor.py +158 -0
- ob_postgres-0.1.0/src/ob_postgres/exceptions.py +32 -0
- ob_postgres-0.1.0/src/ob_postgres/type_codes.py +45 -0
- ob_postgres-0.1.0/tests/test_driver.py +358 -0
- ob_postgres-0.1.0/tests/test_yaml_detection.py +29 -0
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
__pycache__/
|
|
2
|
+
*.py[cod]
|
|
3
|
+
*$py.class
|
|
4
|
+
*.so
|
|
5
|
+
dist/
|
|
6
|
+
build/
|
|
7
|
+
*.egg-info/
|
|
8
|
+
*.egg
|
|
9
|
+
.eggs/
|
|
10
|
+
.venv/
|
|
11
|
+
venv/
|
|
12
|
+
.env
|
|
13
|
+
*.db
|
|
14
|
+
*.duckdb
|
|
15
|
+
*.sqlite3
|
|
16
|
+
.mypy_cache/
|
|
17
|
+
.ruff_cache/
|
|
18
|
+
.pytest_cache/
|
|
19
|
+
htmlcov/
|
|
20
|
+
.coverage
|
|
21
|
+
.coverage.*
|
|
22
|
+
*.log
|
|
23
|
+
.DS_Store
|
|
24
|
+
.claude/
|
|
25
|
+
security/
|
|
26
|
+
infra/
|
|
27
|
+
site/
|
|
28
|
+
scripts/
|
|
29
|
+
.mcpregistry_github_token
|
|
30
|
+
.mcpregistry_registry_token
|
|
31
|
+
_*/
|
|
32
|
+
blog/
|
|
33
|
+
review*
|
|
34
|
+
design/
|
|
35
|
+
examples/test*
|
|
@@ -0,0 +1,109 @@
|
|
|
1
|
+
# ob-postgres — OrionBelt Semantic Layer Driver for Postgres
|
|
2
|
+
|
|
3
|
+
## Purpose
|
|
4
|
+
|
|
5
|
+
PEP 249 DB-API 2.0 driver wrapping `adbc-driver-postgresql = ">=1.0"` that intercepts OBML YAML
|
|
6
|
+
queries, compiles them to postgres SQL via the OrionBelt CompilationPipeline
|
|
7
|
+
(direct import) or OB REST API (standalone), and executes them natively.
|
|
8
|
+
|
|
9
|
+
Uses ADBC (Arrow Database Connectivity) for native Arrow support — `fetch_arrow_table()`
|
|
10
|
+
returns zero-copy PyArrow Tables directly from the PostgreSQL wire protocol.
|
|
11
|
+
|
|
12
|
+
**OB dialect string:** `"postgres"`
|
|
13
|
+
**Author:** Ralf Becher / RALFORION d.o.o. (info@orionbelt.ai)
|
|
14
|
+
**License:** Apache 2.0
|
|
15
|
+
|
|
16
|
+
---
|
|
17
|
+
|
|
18
|
+
## Module Map
|
|
19
|
+
|
|
20
|
+
| File | Responsibility |
|
|
21
|
+
|---|---|
|
|
22
|
+
| `__init__.py` | PEP 249 connect() + module constants |
|
|
23
|
+
| `connection.py` | Connection class |
|
|
24
|
+
| `cursor.py` | Cursor class — OBML detection + execution + Arrow |
|
|
25
|
+
| `compiler.py` | Direct OB import or REST fallback |
|
|
26
|
+
| `exceptions.py` | PEP 249 exception hierarchy |
|
|
27
|
+
| `type_codes.py` | PEP 249 type objects |
|
|
28
|
+
|
|
29
|
+
---
|
|
30
|
+
|
|
31
|
+
## connect() Parameters
|
|
32
|
+
|
|
33
|
+
### postgres-specific
|
|
34
|
+
| host | str | PostgreSQL host (default: localhost) |
|
|
35
|
+
| port | int | PostgreSQL port (default: 5432) |
|
|
36
|
+
| dbname | str | Database name |
|
|
37
|
+
| user | str | Username |
|
|
38
|
+
| password | str | Password |
|
|
39
|
+
| sslmode | str | SSL mode (disable/require/verify-full) |
|
|
40
|
+
|
|
41
|
+
### OrionBelt parameters (same across all vendors)
|
|
42
|
+
| Parameter | Default | Description |
|
|
43
|
+
|---|---|---|
|
|
44
|
+
| ob_model_id | None | Pre-loaded OB model ID |
|
|
45
|
+
| ob_model_yaml | None | OBML YAML string to load |
|
|
46
|
+
| ob_model_file | None | Path to .obml.yaml file |
|
|
47
|
+
| ob_api_url | http://localhost:8000 | OB REST API URL (standalone mode only) |
|
|
48
|
+
| ob_timeout | 30 | HTTP timeout in seconds |
|
|
49
|
+
|
|
50
|
+
---
|
|
51
|
+
|
|
52
|
+
## Compiler Bridge (compiler.py)
|
|
53
|
+
|
|
54
|
+
Auto-detects whether OB core is importable (same process as orionbelt-api)
|
|
55
|
+
or whether to fall back to REST:
|
|
56
|
+
|
|
57
|
+
```python
|
|
58
|
+
def compile_obml(obml: dict, model, dialect: str) -> str:
|
|
59
|
+
try:
|
|
60
|
+
# Direct call — zero overhead, used when embedded in orionbelt-api
|
|
61
|
+
from orionbelt.compiler.pipeline import CompilationPipeline
|
|
62
|
+
result = CompilationPipeline().compile(obml, model, dialect)
|
|
63
|
+
return result.sql
|
|
64
|
+
except ImportError:
|
|
65
|
+
# Standalone mode — call REST API
|
|
66
|
+
return _compile_via_rest(obml, dialect)
|
|
67
|
+
```
|
|
68
|
+
|
|
69
|
+
---
|
|
70
|
+
|
|
71
|
+
## Vendor-Specific Notes
|
|
72
|
+
|
|
73
|
+
- ADBC uses ? (qmark) paramstyle; set paramstyle="qmark" in module globals
|
|
74
|
+
- connect() builds a PostgreSQL URI from keyword args for ADBC
|
|
75
|
+
- commit() and rollback() are meaningful — Postgres uses real transactions
|
|
76
|
+
- fetch_arrow_table() returns zero-copy PyArrow Table via ADBC native Arrow support
|
|
77
|
+
- orionbelt_1 schema validated against this driver; fan-trap prevention via UNION ALL is critical
|
|
78
|
+
|
|
79
|
+
---
|
|
80
|
+
|
|
81
|
+
## Type System
|
|
82
|
+
|
|
83
|
+
ADBC cursor description may return OIDs or ADBC type identifiers.
|
|
84
|
+
PG_OID_MAP maps common OIDs: 23=INT4, 25=TEXT, 700=FLOAT4, 701=FLOAT8, 1114=TIMESTAMP.
|
|
85
|
+
Arrow path bypasses type_code mapping entirely — uses Arrow schema types directly.
|
|
86
|
+
|
|
87
|
+
---
|
|
88
|
+
|
|
89
|
+
## Build Order
|
|
90
|
+
|
|
91
|
+
Session 1: exceptions.py + type_codes.py + compiler.py + unit tests
|
|
92
|
+
Session 2: connection.py + cursor.py (OBML detection, core execute)
|
|
93
|
+
Session 3: __init__.py + connect() + integration tests
|
|
94
|
+
Session 4: SQLAlchemy dialect (ob+postgres:// URL scheme)
|
|
95
|
+
|
|
96
|
+
---
|
|
97
|
+
|
|
98
|
+
## Dependencies
|
|
99
|
+
|
|
100
|
+
```toml
|
|
101
|
+
[project.dependencies]
|
|
102
|
+
adbc-driver-postgresql = ">=1.0"
|
|
103
|
+
pyarrow = ">=16.0"
|
|
104
|
+
pyyaml = ">=6.0"
|
|
105
|
+
|
|
106
|
+
[project.optional-dependencies]
|
|
107
|
+
sqlalchemy = ["sqlalchemy>=2.0"]
|
|
108
|
+
dev = ["pytest>=8.0", "pytest-asyncio>=0.23", "mypy>=1.10", "ruff>=0.4", "respx>=0.21"]
|
|
109
|
+
```
|
|
@@ -0,0 +1,70 @@
|
|
|
1
|
+
Business Source License 1.1
|
|
2
|
+
|
|
3
|
+
Parameters
|
|
4
|
+
|
|
5
|
+
Licensor: RALFORION d.o.o.
|
|
6
|
+
Licensed Work: OrionBelt Semantic Layer 1.0
|
|
7
|
+
The Licensed Work is (c) 2025 RALFORION d.o.o.
|
|
8
|
+
Additional Use Grant: You may use the Licensed Work for any purpose,
|
|
9
|
+
including production use, PROVIDED THAT you may
|
|
10
|
+
not embed OrionBelt in a commercial product or
|
|
11
|
+
provide it as part of a commercial service
|
|
12
|
+
without purchasing a commercial license from
|
|
13
|
+
the Licensor. Contact licensing@ralforion.com
|
|
14
|
+
for commercial licensing inquiries.
|
|
15
|
+
Change Date: 2030-03-16
|
|
16
|
+
Change License: Apache License, Version 2.0
|
|
17
|
+
|
|
18
|
+
-----------------------------------------------------------------------------
|
|
19
|
+
|
|
20
|
+
Business Source License 1.1
|
|
21
|
+
|
|
22
|
+
License text copyright (c) 2017 MariaDB Corporation Ab, All Rights Reserved.
|
|
23
|
+
"Business Source License" is a trademark of MariaDB Corporation Ab.
|
|
24
|
+
|
|
25
|
+
Terms
|
|
26
|
+
|
|
27
|
+
The Licensor hereby grants you the right to copy, modify, create derivative
|
|
28
|
+
works, redistribute, and make non-production use of the Licensed Work. The
|
|
29
|
+
Licensor may make an Additional Use Grant, above, permitting limited
|
|
30
|
+
production use.
|
|
31
|
+
|
|
32
|
+
Effective on the Change Date, or the fourth anniversary of the first publicly
|
|
33
|
+
available distribution of a specific version of the Licensed Work under this
|
|
34
|
+
License, whichever comes first, the Licensor hereby grants you rights under
|
|
35
|
+
the terms of the Change License, and the rights granted in the paragraph
|
|
36
|
+
above terminate.
|
|
37
|
+
|
|
38
|
+
If your use of the Licensed Work does not comply with the requirements
|
|
39
|
+
currently in effect as described in this License, you must purchase a
|
|
40
|
+
commercial license from the Licensor, its affiliated entities, or authorized
|
|
41
|
+
resellers, or you must refrain from using the Licensed Work.
|
|
42
|
+
|
|
43
|
+
All copies of the original and modified Licensed Work, and derivative works
|
|
44
|
+
of the Licensed Work, are subject to this License. This License applies
|
|
45
|
+
separately for each version of the Licensed Work and the Change Date may vary
|
|
46
|
+
for each version of the Licensed Work released by Licensor.
|
|
47
|
+
|
|
48
|
+
You must conspicuously display this License on each original or modified copy
|
|
49
|
+
of the Licensed Work. If you receive the Licensed Work in original or
|
|
50
|
+
modified form from a third party, the terms and conditions set forth in this
|
|
51
|
+
License apply to your use of that work.
|
|
52
|
+
|
|
53
|
+
Any use of the Licensed Work in violation of this License will automatically
|
|
54
|
+
terminate your rights under this License for the current and all other
|
|
55
|
+
versions of the Licensed Work.
|
|
56
|
+
|
|
57
|
+
This License does not grant you any right in any trademark or logo of
|
|
58
|
+
Licensor or its affiliates (provided that you may use a trademark or logo of
|
|
59
|
+
Licensor as expressly required by this License).
|
|
60
|
+
|
|
61
|
+
TO THE EXTENT PERMITTED BY APPLICABLE LAW, THE LICENSED WORK IS PROVIDED ON
|
|
62
|
+
AN "AS IS" BASIS. LICENSOR HEREBY DISCLAIMS ALL WARRANTIES AND CONDITIONS,
|
|
63
|
+
EXPRESS OR IMPLIED, INCLUDING (WITHOUT LIMITATION) WARRANTIES OF
|
|
64
|
+
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, NON-INFRINGEMENT, AND
|
|
65
|
+
TITLE.
|
|
66
|
+
|
|
67
|
+
MariaDB hereby grants you permission to use this License's text to license
|
|
68
|
+
your works, and to refer to it using the trademark "Business Source License",
|
|
69
|
+
as long as you comply with the Covenants of Licensor of MariaDB's Business Source
|
|
70
|
+
License Trademark Policy available at https://mariadb.com/bsl11/.
|
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: ob-postgres
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: OrionBelt Semantic Layer driver for PostgreSQL (DB-API 2.0)
|
|
5
|
+
Author-email: Ralf Becher <info@orionbelt.ai>
|
|
6
|
+
License: BSL-1.1
|
|
7
|
+
License-File: LICENSE
|
|
8
|
+
Keywords: db-api,obml,orionbelt,postgres,semantic-layer
|
|
9
|
+
Requires-Python: >=3.12
|
|
10
|
+
Requires-Dist: adbc-driver-postgresql>=1.0
|
|
11
|
+
Requires-Dist: ob-driver-core>=0.1
|
|
12
|
+
Requires-Dist: pyarrow>=16.0
|
|
13
|
+
Provides-Extra: dev
|
|
14
|
+
Requires-Dist: mypy>=1.10; extra == 'dev'
|
|
15
|
+
Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
|
|
16
|
+
Requires-Dist: pytest>=8.0; extra == 'dev'
|
|
17
|
+
Requires-Dist: respx>=0.21; extra == 'dev'
|
|
18
|
+
Requires-Dist: ruff>=0.4; extra == 'dev'
|
|
19
|
+
Provides-Extra: sqlalchemy
|
|
20
|
+
Requires-Dist: sqlalchemy>=2.0; extra == 'sqlalchemy'
|
|
21
|
+
Description-Content-Type: text/markdown
|
|
22
|
+
|
|
23
|
+
# ob-postgres
|
|
24
|
+
|
|
25
|
+
OrionBelt Semantic Layer driver for PostgreSQL (DB-API 2.0)
|
|
26
|
+
|
|
27
|
+
Part of the [OrionBelt Semantic Layer](https://github.com/ralfbecher/orionbelt-semantic-layer) project.
|
|
28
|
+
|
|
29
|
+
## Installation
|
|
30
|
+
|
|
31
|
+
```bash
|
|
32
|
+
pip install ob-postgres
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
## License
|
|
36
|
+
|
|
37
|
+
BSL-1.1 — see [LICENSE](LICENSE) for details.
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
# ob-postgres
|
|
2
|
+
|
|
3
|
+
OrionBelt Semantic Layer driver for PostgreSQL (DB-API 2.0)
|
|
4
|
+
|
|
5
|
+
Part of the [OrionBelt Semantic Layer](https://github.com/ralfbecher/orionbelt-semantic-layer) project.
|
|
6
|
+
|
|
7
|
+
## Installation
|
|
8
|
+
|
|
9
|
+
```bash
|
|
10
|
+
pip install ob-postgres
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
## License
|
|
14
|
+
|
|
15
|
+
BSL-1.1 — see [LICENSE](LICENSE) for details.
|
|
@@ -0,0 +1,52 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "ob-postgres"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "OrionBelt Semantic Layer driver for PostgreSQL (DB-API 2.0)"
|
|
9
|
+
license = { text = "BSL-1.1" }
|
|
10
|
+
readme = "README.md"
|
|
11
|
+
authors = [{ name = "Ralf Becher", email = "info@orionbelt.ai" }]
|
|
12
|
+
keywords = ["postgres", "semantic-layer", "obml", "db-api", "orionbelt"]
|
|
13
|
+
requires-python = ">=3.12"
|
|
14
|
+
dependencies = [
|
|
15
|
+
"ob-driver-core>=0.1",
|
|
16
|
+
"adbc-driver-postgresql>=1.0",
|
|
17
|
+
"pyarrow>=16.0",
|
|
18
|
+
]
|
|
19
|
+
|
|
20
|
+
[project.optional-dependencies]
|
|
21
|
+
sqlalchemy = ["sqlalchemy>=2.0"]
|
|
22
|
+
dev = [
|
|
23
|
+
"pytest>=8.0",
|
|
24
|
+
"pytest-asyncio>=0.23",
|
|
25
|
+
"mypy>=1.10",
|
|
26
|
+
"ruff>=0.4",
|
|
27
|
+
"respx>=0.21",
|
|
28
|
+
]
|
|
29
|
+
|
|
30
|
+
[project.entry-points."sqlalchemy.dialects"]
|
|
31
|
+
"ob.postgres" = "ob_postgres.dialect:OBPostgresDialect"
|
|
32
|
+
|
|
33
|
+
[tool.hatch.build.targets.wheel]
|
|
34
|
+
packages = ["src/ob_postgres"]
|
|
35
|
+
|
|
36
|
+
[tool.ruff]
|
|
37
|
+
line-length = 100
|
|
38
|
+
target-version = "py312"
|
|
39
|
+
|
|
40
|
+
[tool.ruff.lint]
|
|
41
|
+
select = ["E", "F", "I", "UP", "B", "SIM", "ANN"]
|
|
42
|
+
|
|
43
|
+
[tool.mypy]
|
|
44
|
+
python_version = "3.12"
|
|
45
|
+
strict = true
|
|
46
|
+
files = ["src/ob_postgres"]
|
|
47
|
+
|
|
48
|
+
[tool.pytest.ini_options]
|
|
49
|
+
asyncio_mode = "auto"
|
|
50
|
+
testpaths = ["tests"]
|
|
51
|
+
markers = ["integration: requires live DB + OB API"]
|
|
52
|
+
addopts = "-m 'not integration'"
|
|
@@ -0,0 +1,157 @@
|
|
|
1
|
+
"""ob-postgres — OrionBelt Semantic Layer driver for PostgreSQL (PEP 249 DB-API 2.0).
|
|
2
|
+
|
|
3
|
+
Uses ``adbc-driver-postgresql`` for native Arrow support. ADBC provides
|
|
4
|
+
zero-copy Arrow results directly from the PostgreSQL wire protocol, enabling
|
|
5
|
+
efficient Arrow Flight SQL streaming.
|
|
6
|
+
|
|
7
|
+
Requires the OrionBelt REST API running in single-model mode (MODEL_FILE set).
|
|
8
|
+
OBML queries are compiled to SQL via ``POST /v1/query/sql``.
|
|
9
|
+
|
|
10
|
+
Usage::
|
|
11
|
+
|
|
12
|
+
import ob_postgres
|
|
13
|
+
|
|
14
|
+
conn = ob_postgres.connect(dbname="mydb", user="me", password="secret")
|
|
15
|
+
with conn.cursor() as cur:
|
|
16
|
+
cur.execute("select:\\n dimensions:\\n - Region\\n measures:\\n - Revenue")
|
|
17
|
+
print(cur.fetchall())
|
|
18
|
+
"""
|
|
19
|
+
|
|
20
|
+
from __future__ import annotations
|
|
21
|
+
|
|
22
|
+
from typing import Any
|
|
23
|
+
from urllib.parse import quote_plus, urlencode
|
|
24
|
+
|
|
25
|
+
import adbc_driver_postgresql.dbapi
|
|
26
|
+
|
|
27
|
+
from ob_postgres.connection import Connection
|
|
28
|
+
from ob_postgres.exceptions import (
|
|
29
|
+
DatabaseError,
|
|
30
|
+
DataError,
|
|
31
|
+
Error,
|
|
32
|
+
IntegrityError,
|
|
33
|
+
InterfaceError,
|
|
34
|
+
InternalError,
|
|
35
|
+
NotSupportedError,
|
|
36
|
+
OperationalError,
|
|
37
|
+
ProgrammingError,
|
|
38
|
+
Warning,
|
|
39
|
+
)
|
|
40
|
+
|
|
41
|
+
# PEP 249 module-level constants
|
|
42
|
+
apilevel = "2.0"
|
|
43
|
+
threadsafety = 1 # threads may share the module but not connections
|
|
44
|
+
paramstyle = "qmark" # ADBC uses ? placeholders
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
def _build_pg_uri(
|
|
48
|
+
*,
|
|
49
|
+
host: str,
|
|
50
|
+
port: int,
|
|
51
|
+
dbname: str,
|
|
52
|
+
user: str | None,
|
|
53
|
+
password: str | None,
|
|
54
|
+
sslmode: str | None,
|
|
55
|
+
) -> str:
|
|
56
|
+
"""Build a PostgreSQL connection URI from keyword arguments."""
|
|
57
|
+
userinfo = ""
|
|
58
|
+
if user is not None:
|
|
59
|
+
userinfo = quote_plus(user)
|
|
60
|
+
if password is not None:
|
|
61
|
+
userinfo += f":{quote_plus(password)}"
|
|
62
|
+
userinfo += "@"
|
|
63
|
+
params: dict[str, str] = {}
|
|
64
|
+
if sslmode is not None:
|
|
65
|
+
params["sslmode"] = sslmode
|
|
66
|
+
query = f"?{urlencode(params)}" if params else ""
|
|
67
|
+
return f"postgresql://{userinfo}{host}:{port}/{dbname}{query}"
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
def connect(
|
|
71
|
+
dsn: str | None = None,
|
|
72
|
+
*,
|
|
73
|
+
host: str = "localhost",
|
|
74
|
+
port: int = 5432,
|
|
75
|
+
dbname: str = "postgres",
|
|
76
|
+
user: str | None = None,
|
|
77
|
+
password: str | None = None,
|
|
78
|
+
sslmode: str | None = None,
|
|
79
|
+
schema: str | None = None,
|
|
80
|
+
options: dict[str, Any] | None = None,
|
|
81
|
+
# OrionBelt parameters
|
|
82
|
+
ob_api_url: str = "http://localhost:8000",
|
|
83
|
+
ob_timeout: int = 30,
|
|
84
|
+
) -> Connection:
|
|
85
|
+
"""Open a PostgreSQL connection with OBML support (ADBC driver).
|
|
86
|
+
|
|
87
|
+
Parameters
|
|
88
|
+
----------
|
|
89
|
+
dsn : str, optional
|
|
90
|
+
Full PostgreSQL URI (overrides individual params).
|
|
91
|
+
host : str
|
|
92
|
+
PostgreSQL host (default: ``localhost``).
|
|
93
|
+
port : int
|
|
94
|
+
PostgreSQL port (default: ``5432``).
|
|
95
|
+
dbname : str
|
|
96
|
+
Database name (default: ``postgres``).
|
|
97
|
+
user : str, optional
|
|
98
|
+
Username.
|
|
99
|
+
password : str, optional
|
|
100
|
+
Password.
|
|
101
|
+
sslmode : str, optional
|
|
102
|
+
SSL mode (``disable``, ``require``, ``verify-full``, etc.).
|
|
103
|
+
schema : str, optional
|
|
104
|
+
PostgreSQL schema — sets ``search_path`` after connecting.
|
|
105
|
+
options : dict, optional
|
|
106
|
+
Extra connection options (added as URI query parameters).
|
|
107
|
+
ob_api_url : str
|
|
108
|
+
OrionBelt REST API URL (must be running in single-model mode).
|
|
109
|
+
ob_timeout : int
|
|
110
|
+
HTTP timeout in seconds for OBML compilation.
|
|
111
|
+
"""
|
|
112
|
+
if dsn is not None:
|
|
113
|
+
uri = dsn
|
|
114
|
+
else:
|
|
115
|
+
# Include search_path as a connection option if schema is specified
|
|
116
|
+
extra_options = dict(options) if options else {}
|
|
117
|
+
if schema:
|
|
118
|
+
extra_options["options"] = f"-csearch_path={schema}"
|
|
119
|
+
|
|
120
|
+
uri = _build_pg_uri(
|
|
121
|
+
host=host,
|
|
122
|
+
port=port,
|
|
123
|
+
dbname=dbname,
|
|
124
|
+
user=user,
|
|
125
|
+
password=password,
|
|
126
|
+
sslmode=sslmode,
|
|
127
|
+
)
|
|
128
|
+
if extra_options:
|
|
129
|
+
sep = "&" if "?" in uri else "?"
|
|
130
|
+
uri += sep + urlencode(extra_options)
|
|
131
|
+
|
|
132
|
+
native = adbc_driver_postgresql.dbapi.connect(uri)
|
|
133
|
+
|
|
134
|
+
return Connection(
|
|
135
|
+
native,
|
|
136
|
+
ob_api_url=ob_api_url,
|
|
137
|
+
ob_timeout=ob_timeout,
|
|
138
|
+
)
|
|
139
|
+
|
|
140
|
+
|
|
141
|
+
__all__ = [
|
|
142
|
+
"apilevel",
|
|
143
|
+
"threadsafety",
|
|
144
|
+
"paramstyle",
|
|
145
|
+
"connect",
|
|
146
|
+
"Connection",
|
|
147
|
+
"Warning",
|
|
148
|
+
"Error",
|
|
149
|
+
"InterfaceError",
|
|
150
|
+
"DatabaseError",
|
|
151
|
+
"DataError",
|
|
152
|
+
"OperationalError",
|
|
153
|
+
"IntegrityError",
|
|
154
|
+
"InternalError",
|
|
155
|
+
"ProgrammingError",
|
|
156
|
+
"NotSupportedError",
|
|
157
|
+
]
|
|
@@ -0,0 +1,8 @@
|
|
|
1
|
+
"""OrionBelt compilation bridge — delegates to ob-driver-core."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from ob_driver_core.compiler import compile_obml
|
|
6
|
+
from ob_driver_core.detection import is_obml, parse_obml
|
|
7
|
+
|
|
8
|
+
__all__ = ["compile_obml", "is_obml", "parse_obml"]
|
|
@@ -0,0 +1,65 @@
|
|
|
1
|
+
"""PEP 249 Connection wrapping ``adbc_driver_postgresql.dbapi.Connection``."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from typing import Any
|
|
6
|
+
|
|
7
|
+
from ob_postgres.cursor import Cursor
|
|
8
|
+
from ob_postgres.exceptions import ProgrammingError
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class Connection:
|
|
12
|
+
"""DB-API 2.0 connection that wraps an ADBC PostgreSQL connection.
|
|
13
|
+
|
|
14
|
+
ADBC provides native Arrow support — cursors created from this connection
|
|
15
|
+
support ``fetch_arrow_table()`` for zero-copy Arrow Flight SQL streaming.
|
|
16
|
+
OBML queries are compiled to SQL via the OrionBelt REST API.
|
|
17
|
+
"""
|
|
18
|
+
|
|
19
|
+
def __init__(
|
|
20
|
+
self,
|
|
21
|
+
native: Any,
|
|
22
|
+
*,
|
|
23
|
+
ob_api_url: str = "http://localhost:8000",
|
|
24
|
+
ob_timeout: int = 30,
|
|
25
|
+
) -> None:
|
|
26
|
+
self._native = native
|
|
27
|
+
self._closed = False
|
|
28
|
+
self._ob_api_url = ob_api_url
|
|
29
|
+
self._ob_timeout = ob_timeout
|
|
30
|
+
|
|
31
|
+
def _check_open(self) -> None:
|
|
32
|
+
if self._closed:
|
|
33
|
+
raise ProgrammingError("Connection is closed.")
|
|
34
|
+
|
|
35
|
+
def cursor(self) -> Cursor:
|
|
36
|
+
"""Return a new Cursor for this connection."""
|
|
37
|
+
self._check_open()
|
|
38
|
+
native_cursor = self._native.cursor()
|
|
39
|
+
return Cursor(
|
|
40
|
+
native_cursor,
|
|
41
|
+
ob_api_url=self._ob_api_url,
|
|
42
|
+
ob_timeout=self._ob_timeout,
|
|
43
|
+
)
|
|
44
|
+
|
|
45
|
+
def commit(self) -> None:
|
|
46
|
+
"""Commit the current transaction."""
|
|
47
|
+
self._check_open()
|
|
48
|
+
self._native.commit()
|
|
49
|
+
|
|
50
|
+
def rollback(self) -> None:
|
|
51
|
+
"""Rollback the current transaction."""
|
|
52
|
+
self._check_open()
|
|
53
|
+
self._native.rollback()
|
|
54
|
+
|
|
55
|
+
def close(self) -> None:
|
|
56
|
+
"""Close the connection."""
|
|
57
|
+
if not self._closed:
|
|
58
|
+
self._native.close()
|
|
59
|
+
self._closed = True
|
|
60
|
+
|
|
61
|
+
def __enter__(self) -> Connection:
|
|
62
|
+
return self
|
|
63
|
+
|
|
64
|
+
def __exit__(self, *exc: object) -> None:
|
|
65
|
+
self.close()
|
|
@@ -0,0 +1,158 @@
|
|
|
1
|
+
"""PEP 249 Cursor with OBML interception for PostgreSQL.
|
|
2
|
+
|
|
3
|
+
Uses ``adbc-driver-postgresql`` (Arrow Database Connectivity) for native
|
|
4
|
+
Arrow support. ``fetch_arrow_table()`` returns a PyArrow Table directly
|
|
5
|
+
from the ADBC cursor — zero-copy columnar transfer for Arrow Flight SQL.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
from typing import Any
|
|
11
|
+
|
|
12
|
+
from ob_postgres.compiler import compile_obml, is_obml, parse_obml
|
|
13
|
+
from ob_postgres.exceptions import NotSupportedError, ProgrammingError
|
|
14
|
+
from ob_postgres.type_codes import PG_OID_MAP, STRING
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class Cursor:
|
|
18
|
+
"""DB-API 2.0 cursor that intercepts OBML YAML queries.
|
|
19
|
+
|
|
20
|
+
If the query is OBML, it is compiled to PostgreSQL SQL via the OrionBelt
|
|
21
|
+
REST API before execution. Plain SQL is passed through unchanged.
|
|
22
|
+
|
|
23
|
+
Wraps an ``adbc_driver_postgresql.dbapi`` cursor which natively supports
|
|
24
|
+
``fetch_arrow_table()`` for zero-copy Arrow results.
|
|
25
|
+
"""
|
|
26
|
+
|
|
27
|
+
arraysize: int = 1
|
|
28
|
+
|
|
29
|
+
def __init__(
|
|
30
|
+
self,
|
|
31
|
+
native: Any,
|
|
32
|
+
*,
|
|
33
|
+
ob_api_url: str = "http://localhost:8000",
|
|
34
|
+
ob_timeout: int = 30,
|
|
35
|
+
) -> None:
|
|
36
|
+
self._native = native
|
|
37
|
+
self._closed = False
|
|
38
|
+
self._ob_api_url = ob_api_url
|
|
39
|
+
self._ob_timeout = ob_timeout
|
|
40
|
+
|
|
41
|
+
@property
|
|
42
|
+
def description(
|
|
43
|
+
self,
|
|
44
|
+
) -> tuple[tuple[str, Any, None, None, None, None, None], ...] | None:
|
|
45
|
+
"""PEP 249 cursor description — 7-item tuples per column."""
|
|
46
|
+
native_desc = self._native.description
|
|
47
|
+
if native_desc is None:
|
|
48
|
+
return None
|
|
49
|
+
cols: list[tuple[str, Any, None, None, None, None, None]] = []
|
|
50
|
+
for col in native_desc:
|
|
51
|
+
name = col[0]
|
|
52
|
+
type_code = col[1]
|
|
53
|
+
# ADBC PG driver may return OIDs or other type identifiers
|
|
54
|
+
mapped = PG_OID_MAP.get(type_code, STRING) if isinstance(type_code, int) else STRING
|
|
55
|
+
cols.append((name, mapped, None, None, None, None, None))
|
|
56
|
+
return tuple(cols)
|
|
57
|
+
|
|
58
|
+
@property
|
|
59
|
+
def rowcount(self) -> int:
|
|
60
|
+
return self._native.rowcount
|
|
61
|
+
|
|
62
|
+
def _check_open(self) -> None:
|
|
63
|
+
if self._closed:
|
|
64
|
+
raise ProgrammingError("Cursor is closed.")
|
|
65
|
+
|
|
66
|
+
def execute(self, operation: str, parameters: Any = None) -> Cursor:
|
|
67
|
+
"""Execute a query — OBML YAML or plain SQL."""
|
|
68
|
+
self._check_open()
|
|
69
|
+
sql = self._resolve_sql(operation)
|
|
70
|
+
if parameters is not None:
|
|
71
|
+
self._native.execute(sql, parameters)
|
|
72
|
+
else:
|
|
73
|
+
self._native.execute(sql)
|
|
74
|
+
return self
|
|
75
|
+
|
|
76
|
+
def executemany(self, operation: str, seq_of_parameters: Any) -> None:
|
|
77
|
+
"""Execute against all parameter sequences.
|
|
78
|
+
|
|
79
|
+
OBML queries are not supported with executemany — raises NotSupportedError.
|
|
80
|
+
"""
|
|
81
|
+
self._check_open()
|
|
82
|
+
if is_obml(operation):
|
|
83
|
+
raise NotSupportedError("executemany() is not supported for OBML queries.")
|
|
84
|
+
for params in seq_of_parameters:
|
|
85
|
+
self._native.execute(operation, params)
|
|
86
|
+
|
|
87
|
+
def _resolve_sql(self, operation: str) -> str:
|
|
88
|
+
"""Compile OBML to SQL or return plain SQL unchanged."""
|
|
89
|
+
if not is_obml(operation):
|
|
90
|
+
return operation
|
|
91
|
+
obml = parse_obml(operation)
|
|
92
|
+
return compile_obml(
|
|
93
|
+
obml,
|
|
94
|
+
dialect="postgres",
|
|
95
|
+
ob_api_url=self._ob_api_url,
|
|
96
|
+
ob_timeout=self._ob_timeout,
|
|
97
|
+
)
|
|
98
|
+
|
|
99
|
+
def fetchone(self) -> tuple[Any, ...] | None:
|
|
100
|
+
"""Fetch the next row."""
|
|
101
|
+
self._check_open()
|
|
102
|
+
row = self._native.fetchone()
|
|
103
|
+
return tuple(row) if row is not None else None
|
|
104
|
+
|
|
105
|
+
def fetchmany(self, size: int | None = None) -> list[tuple[Any, ...]]:
|
|
106
|
+
"""Fetch the next *size* rows."""
|
|
107
|
+
self._check_open()
|
|
108
|
+
n = size if size is not None else self.arraysize
|
|
109
|
+
rows = self._native.fetchmany(n)
|
|
110
|
+
return [tuple(r) for r in rows]
|
|
111
|
+
|
|
112
|
+
def fetchall(self) -> list[tuple[Any, ...]]:
|
|
113
|
+
"""Fetch all remaining rows."""
|
|
114
|
+
self._check_open()
|
|
115
|
+
rows = self._native.fetchall()
|
|
116
|
+
return [tuple(r) for r in rows]
|
|
117
|
+
|
|
118
|
+
def fetch_arrow_table(self) -> Any:
|
|
119
|
+
"""Fetch all remaining rows as a PyArrow Table (ADBC native).
|
|
120
|
+
|
|
121
|
+
ADBC provides zero-copy Arrow results directly from the PostgreSQL
|
|
122
|
+
wire protocol. Significantly more memory-efficient than materialising
|
|
123
|
+
Python row objects and enables efficient Arrow Flight SQL streaming.
|
|
124
|
+
"""
|
|
125
|
+
self._check_open()
|
|
126
|
+
return self._native.fetch_arrow_table()
|
|
127
|
+
|
|
128
|
+
def close(self) -> None:
|
|
129
|
+
"""Close the cursor."""
|
|
130
|
+
if not self._closed:
|
|
131
|
+
self._native.close()
|
|
132
|
+
self._closed = True
|
|
133
|
+
|
|
134
|
+
def setinputsizes(self, sizes: Any) -> None:
|
|
135
|
+
"""No-op — required by PEP 249."""
|
|
136
|
+
|
|
137
|
+
def setoutputsize(self, size: int, column: int | None = None) -> None:
|
|
138
|
+
"""No-op — required by PEP 249."""
|
|
139
|
+
|
|
140
|
+
@property
|
|
141
|
+
def lastrowid(self) -> None:
|
|
142
|
+
"""PostgreSQL does not expose lastrowid via ADBC cursor."""
|
|
143
|
+
return None
|
|
144
|
+
|
|
145
|
+
def __enter__(self) -> Cursor:
|
|
146
|
+
return self
|
|
147
|
+
|
|
148
|
+
def __exit__(self, *exc: object) -> None:
|
|
149
|
+
self.close()
|
|
150
|
+
|
|
151
|
+
def __iter__(self) -> Cursor:
|
|
152
|
+
return self
|
|
153
|
+
|
|
154
|
+
def __next__(self) -> tuple[Any, ...]:
|
|
155
|
+
row = self.fetchone()
|
|
156
|
+
if row is None:
|
|
157
|
+
raise StopIteration
|
|
158
|
+
return row
|
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
"""PEP 249 exceptions — re-exported from ob-driver-core for convenience.
|
|
2
|
+
|
|
3
|
+
Allows users to do: ``from ob_postgres.exceptions import ProgrammingError``
|
|
4
|
+
"""
|
|
5
|
+
|
|
6
|
+
from __future__ import annotations
|
|
7
|
+
|
|
8
|
+
from ob_driver_core.exceptions import (
|
|
9
|
+
DataError,
|
|
10
|
+
DatabaseError,
|
|
11
|
+
Error,
|
|
12
|
+
IntegrityError,
|
|
13
|
+
InterfaceError,
|
|
14
|
+
InternalError,
|
|
15
|
+
NotSupportedError,
|
|
16
|
+
OperationalError,
|
|
17
|
+
ProgrammingError,
|
|
18
|
+
Warning,
|
|
19
|
+
)
|
|
20
|
+
|
|
21
|
+
__all__ = [
|
|
22
|
+
"Warning",
|
|
23
|
+
"Error",
|
|
24
|
+
"InterfaceError",
|
|
25
|
+
"DatabaseError",
|
|
26
|
+
"DataError",
|
|
27
|
+
"OperationalError",
|
|
28
|
+
"IntegrityError",
|
|
29
|
+
"InternalError",
|
|
30
|
+
"ProgrammingError",
|
|
31
|
+
"NotSupportedError",
|
|
32
|
+
]
|
|
@@ -0,0 +1,45 @@
|
|
|
1
|
+
"""PEP 249 type objects — re-exported from ob-driver-core plus Postgres OID mapping.
|
|
2
|
+
|
|
3
|
+
psycopg2 cursor.description returns integer OIDs as type_code.
|
|
4
|
+
The PG_OID_MAP maps common OIDs to PEP 249 type objects.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
from ob_driver_core.type_codes import BINARY, DATETIME, NUMBER, ROWID, STRING
|
|
10
|
+
|
|
11
|
+
__all__ = ["STRING", "BINARY", "NUMBER", "DATETIME", "ROWID", "PG_OID_MAP"]
|
|
12
|
+
|
|
13
|
+
# Common PostgreSQL OIDs — see pg_type catalog
|
|
14
|
+
# https://www.postgresql.org/docs/current/catalog-pg-type.html
|
|
15
|
+
PG_OID_MAP: dict[int, object] = {
|
|
16
|
+
# Boolean
|
|
17
|
+
16: STRING, # bool
|
|
18
|
+
# Numeric types
|
|
19
|
+
20: NUMBER, # int8 (bigint)
|
|
20
|
+
21: NUMBER, # int2 (smallint)
|
|
21
|
+
23: NUMBER, # int4 (integer)
|
|
22
|
+
26: NUMBER, # oid
|
|
23
|
+
700: NUMBER, # float4
|
|
24
|
+
701: NUMBER, # float8
|
|
25
|
+
790: NUMBER, # money
|
|
26
|
+
1700: NUMBER, # numeric/decimal
|
|
27
|
+
# String types
|
|
28
|
+
18: STRING, # char
|
|
29
|
+
19: STRING, # name
|
|
30
|
+
25: STRING, # text
|
|
31
|
+
1042: STRING, # bpchar (char(n))
|
|
32
|
+
1043: STRING, # varchar
|
|
33
|
+
2950: STRING, # uuid
|
|
34
|
+
114: STRING, # json
|
|
35
|
+
3802: STRING, # jsonb
|
|
36
|
+
# Date/time types
|
|
37
|
+
1082: DATETIME, # date
|
|
38
|
+
1083: DATETIME, # time
|
|
39
|
+
1114: DATETIME, # timestamp
|
|
40
|
+
1184: DATETIME, # timestamptz
|
|
41
|
+
1186: DATETIME, # interval
|
|
42
|
+
1266: DATETIME, # timetz
|
|
43
|
+
# Binary types
|
|
44
|
+
17: BINARY, # bytea
|
|
45
|
+
}
|
|
@@ -0,0 +1,358 @@
|
|
|
1
|
+
"""Unit tests for the ob-postgres DB-API 2.0 driver.
|
|
2
|
+
|
|
3
|
+
All tests mock psycopg2 — no live PostgreSQL needed.
|
|
4
|
+
OBML tests additionally mock the REST API call to ``/v1/query/sql``.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
from collections import namedtuple
|
|
10
|
+
from unittest.mock import MagicMock, patch
|
|
11
|
+
|
|
12
|
+
import pytest
|
|
13
|
+
|
|
14
|
+
import ob_postgres
|
|
15
|
+
from ob_postgres.connection import Connection
|
|
16
|
+
from ob_postgres.cursor import Cursor
|
|
17
|
+
from ob_postgres.exceptions import NotSupportedError, ProgrammingError
|
|
18
|
+
from ob_postgres.type_codes import DATETIME, NUMBER, STRING
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
# ---------------------------------------------------------------------------
|
|
22
|
+
# Helper to build a mock psycopg2 connection
|
|
23
|
+
# ---------------------------------------------------------------------------
|
|
24
|
+
|
|
25
|
+
# psycopg2 description columns are named tuples
|
|
26
|
+
PgColumn = namedtuple("Column", ["name", "type_code", "display_size",
|
|
27
|
+
"internal_size", "precision", "scale", "null_ok"])
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def _make_mock_native() -> MagicMock:
|
|
31
|
+
"""Return a mock psycopg2 connection."""
|
|
32
|
+
mock_conn = MagicMock()
|
|
33
|
+
mock_cursor = MagicMock()
|
|
34
|
+
mock_conn.cursor.return_value = mock_cursor
|
|
35
|
+
# Default: no results
|
|
36
|
+
mock_cursor.description = None
|
|
37
|
+
mock_cursor.rowcount = -1
|
|
38
|
+
mock_cursor.fetchone.return_value = None
|
|
39
|
+
mock_cursor.fetchall.return_value = []
|
|
40
|
+
mock_cursor.fetchmany.return_value = []
|
|
41
|
+
return mock_conn
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def _mock_api_response(sql: str) -> MagicMock:
|
|
45
|
+
"""Create a mock httpx response returning the given SQL."""
|
|
46
|
+
resp = MagicMock()
|
|
47
|
+
resp.is_success = True
|
|
48
|
+
resp.json.return_value = {"sql": sql}
|
|
49
|
+
return resp
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
# ---------------------------------------------------------------------------
|
|
53
|
+
# PEP 249 module-level constants
|
|
54
|
+
# ---------------------------------------------------------------------------
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def test_apilevel() -> None:
|
|
58
|
+
assert ob_postgres.apilevel == "2.0"
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
def test_threadsafety() -> None:
|
|
62
|
+
assert ob_postgres.threadsafety == 1
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
def test_paramstyle() -> None:
|
|
66
|
+
assert ob_postgres.paramstyle == "format"
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
# ---------------------------------------------------------------------------
|
|
70
|
+
# connect()
|
|
71
|
+
# ---------------------------------------------------------------------------
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
def test_connect_returns_connection() -> None:
|
|
75
|
+
with patch("psycopg2.connect") as mock_connect:
|
|
76
|
+
mock_connect.return_value = _make_mock_native()
|
|
77
|
+
conn = ob_postgres.connect(dbname="testdb", user="testuser")
|
|
78
|
+
assert isinstance(conn, Connection)
|
|
79
|
+
mock_connect.assert_called_once()
|
|
80
|
+
kwargs = mock_connect.call_args.kwargs
|
|
81
|
+
assert kwargs["dbname"] == "testdb"
|
|
82
|
+
assert kwargs["user"] == "testuser"
|
|
83
|
+
|
|
84
|
+
|
|
85
|
+
def test_connect_with_dsn() -> None:
|
|
86
|
+
with patch("psycopg2.connect") as mock_connect:
|
|
87
|
+
mock_connect.return_value = _make_mock_native()
|
|
88
|
+
conn = ob_postgres.connect(dsn="host=localhost dbname=mydb")
|
|
89
|
+
assert isinstance(conn, Connection)
|
|
90
|
+
kwargs = mock_connect.call_args.kwargs
|
|
91
|
+
assert kwargs["dsn"] == "host=localhost dbname=mydb"
|
|
92
|
+
|
|
93
|
+
|
|
94
|
+
def test_connect_context_manager() -> None:
|
|
95
|
+
with patch("psycopg2.connect") as mock_connect:
|
|
96
|
+
mock_native = _make_mock_native()
|
|
97
|
+
mock_connect.return_value = mock_native
|
|
98
|
+
with ob_postgres.connect() as conn:
|
|
99
|
+
assert isinstance(conn, Connection)
|
|
100
|
+
mock_native.close.assert_called_once()
|
|
101
|
+
|
|
102
|
+
|
|
103
|
+
# ---------------------------------------------------------------------------
|
|
104
|
+
# Connection
|
|
105
|
+
# ---------------------------------------------------------------------------
|
|
106
|
+
|
|
107
|
+
|
|
108
|
+
def test_connection_close_is_idempotent() -> None:
|
|
109
|
+
mock_native = _make_mock_native()
|
|
110
|
+
conn = Connection(mock_native)
|
|
111
|
+
conn.close()
|
|
112
|
+
conn.close() # should not raise
|
|
113
|
+
mock_native.close.assert_called_once()
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
def test_connection_cursor_after_close_raises() -> None:
|
|
117
|
+
mock_native = _make_mock_native()
|
|
118
|
+
conn = Connection(mock_native)
|
|
119
|
+
conn.close()
|
|
120
|
+
with pytest.raises(ProgrammingError, match="closed"):
|
|
121
|
+
conn.cursor()
|
|
122
|
+
|
|
123
|
+
|
|
124
|
+
def test_connection_commit() -> None:
|
|
125
|
+
mock_native = _make_mock_native()
|
|
126
|
+
conn = Connection(mock_native)
|
|
127
|
+
conn.commit()
|
|
128
|
+
mock_native.commit.assert_called_once()
|
|
129
|
+
|
|
130
|
+
|
|
131
|
+
def test_connection_rollback() -> None:
|
|
132
|
+
mock_native = _make_mock_native()
|
|
133
|
+
conn = Connection(mock_native)
|
|
134
|
+
conn.rollback()
|
|
135
|
+
mock_native.rollback.assert_called_once()
|
|
136
|
+
|
|
137
|
+
|
|
138
|
+
# ---------------------------------------------------------------------------
|
|
139
|
+
# Cursor — plain SQL
|
|
140
|
+
# ---------------------------------------------------------------------------
|
|
141
|
+
|
|
142
|
+
|
|
143
|
+
def test_cursor_execute_calls_native() -> None:
|
|
144
|
+
mock_native = _make_mock_native()
|
|
145
|
+
conn = Connection(mock_native)
|
|
146
|
+
with conn.cursor() as cur:
|
|
147
|
+
cur.execute("SELECT 1")
|
|
148
|
+
mock_native.cursor().execute.assert_called_once_with("SELECT 1")
|
|
149
|
+
|
|
150
|
+
|
|
151
|
+
def test_cursor_execute_with_params() -> None:
|
|
152
|
+
mock_native = _make_mock_native()
|
|
153
|
+
conn = Connection(mock_native)
|
|
154
|
+
with conn.cursor() as cur:
|
|
155
|
+
cur.execute("SELECT %s + %s", [3, 4])
|
|
156
|
+
mock_native.cursor().execute.assert_called_once_with("SELECT %s + %s", [3, 4])
|
|
157
|
+
|
|
158
|
+
|
|
159
|
+
def test_cursor_execute_returns_self() -> None:
|
|
160
|
+
mock_native = _make_mock_native()
|
|
161
|
+
conn = Connection(mock_native)
|
|
162
|
+
with conn.cursor() as cur:
|
|
163
|
+
result = cur.execute("SELECT 1")
|
|
164
|
+
assert result is cur
|
|
165
|
+
|
|
166
|
+
|
|
167
|
+
def test_cursor_description() -> None:
|
|
168
|
+
mock_native = _make_mock_native()
|
|
169
|
+
mock_cursor = mock_native.cursor()
|
|
170
|
+
mock_cursor.description = [
|
|
171
|
+
PgColumn("num", 23, None, None, None, None, None), # int4 → NUMBER
|
|
172
|
+
PgColumn("txt", 25, None, None, None, None, None), # text → STRING
|
|
173
|
+
PgColumn("dt", 1114, None, None, None, None, None), # timestamp → DATETIME
|
|
174
|
+
]
|
|
175
|
+
conn = Connection(mock_native)
|
|
176
|
+
cur = conn.cursor()
|
|
177
|
+
desc = cur.description
|
|
178
|
+
assert desc is not None
|
|
179
|
+
assert len(desc) == 3
|
|
180
|
+
assert all(len(col) == 7 for col in desc)
|
|
181
|
+
assert desc[0][0] == "num"
|
|
182
|
+
assert desc[0][1] == NUMBER
|
|
183
|
+
assert desc[1][0] == "txt"
|
|
184
|
+
assert desc[1][1] == STRING
|
|
185
|
+
assert desc[2][0] == "dt"
|
|
186
|
+
assert desc[2][1] == DATETIME
|
|
187
|
+
|
|
188
|
+
|
|
189
|
+
def test_cursor_description_none_before_execute() -> None:
|
|
190
|
+
mock_native = _make_mock_native()
|
|
191
|
+
mock_native.cursor().description = None
|
|
192
|
+
conn = Connection(mock_native)
|
|
193
|
+
cur = conn.cursor()
|
|
194
|
+
assert cur.description is None
|
|
195
|
+
|
|
196
|
+
|
|
197
|
+
def test_cursor_fetchone() -> None:
|
|
198
|
+
mock_native = _make_mock_native()
|
|
199
|
+
mock_native.cursor().fetchone.return_value = (42,)
|
|
200
|
+
conn = Connection(mock_native)
|
|
201
|
+
cur = conn.cursor()
|
|
202
|
+
row = cur.fetchone()
|
|
203
|
+
assert row == (42,)
|
|
204
|
+
|
|
205
|
+
|
|
206
|
+
def test_cursor_fetchone_exhausted() -> None:
|
|
207
|
+
mock_native = _make_mock_native()
|
|
208
|
+
mock_native.cursor().fetchone.return_value = None
|
|
209
|
+
conn = Connection(mock_native)
|
|
210
|
+
cur = conn.cursor()
|
|
211
|
+
assert cur.fetchone() is None
|
|
212
|
+
|
|
213
|
+
|
|
214
|
+
def test_cursor_fetchall() -> None:
|
|
215
|
+
mock_native = _make_mock_native()
|
|
216
|
+
mock_native.cursor().fetchall.return_value = [(1, "a"), (2, "b"), (3, "c")]
|
|
217
|
+
conn = Connection(mock_native)
|
|
218
|
+
cur = conn.cursor()
|
|
219
|
+
rows = cur.fetchall()
|
|
220
|
+
assert len(rows) == 3
|
|
221
|
+
assert rows[0] == (1, "a")
|
|
222
|
+
|
|
223
|
+
|
|
224
|
+
def test_cursor_fetchmany() -> None:
|
|
225
|
+
mock_native = _make_mock_native()
|
|
226
|
+
mock_native.cursor().fetchmany.return_value = [(1,), (2,), (3,)]
|
|
227
|
+
conn = Connection(mock_native)
|
|
228
|
+
cur = conn.cursor()
|
|
229
|
+
batch = cur.fetchmany(3)
|
|
230
|
+
assert len(batch) == 3
|
|
231
|
+
mock_native.cursor().fetchmany.assert_called_with(3)
|
|
232
|
+
|
|
233
|
+
|
|
234
|
+
def test_cursor_iteration() -> None:
|
|
235
|
+
mock_native = _make_mock_native()
|
|
236
|
+
mock_native.cursor().fetchone.side_effect = [(0,), (1,), (2,), None]
|
|
237
|
+
conn = Connection(mock_native)
|
|
238
|
+
cur = conn.cursor()
|
|
239
|
+
rows = list(cur)
|
|
240
|
+
assert len(rows) == 3
|
|
241
|
+
assert rows[0] == (0,)
|
|
242
|
+
|
|
243
|
+
|
|
244
|
+
def test_cursor_close_then_fetch_raises() -> None:
|
|
245
|
+
mock_native = _make_mock_native()
|
|
246
|
+
conn = Connection(mock_native)
|
|
247
|
+
cur = conn.cursor()
|
|
248
|
+
cur.close()
|
|
249
|
+
with pytest.raises(ProgrammingError, match="closed"):
|
|
250
|
+
cur.fetchone()
|
|
251
|
+
|
|
252
|
+
|
|
253
|
+
def test_cursor_executemany_plain_sql() -> None:
|
|
254
|
+
mock_native = _make_mock_native()
|
|
255
|
+
conn = Connection(mock_native)
|
|
256
|
+
cur = conn.cursor()
|
|
257
|
+
cur.executemany("INSERT INTO t VALUES (%s, %s)", [(1, "a"), (2, "b")])
|
|
258
|
+
assert mock_native.cursor().execute.call_count == 2
|
|
259
|
+
|
|
260
|
+
|
|
261
|
+
def test_cursor_executemany_obml_raises() -> None:
|
|
262
|
+
mock_native = _make_mock_native()
|
|
263
|
+
conn = Connection(mock_native)
|
|
264
|
+
cur = conn.cursor()
|
|
265
|
+
obml = "select:\n dimensions:\n - Region\n measures:\n - Revenue\n"
|
|
266
|
+
with pytest.raises(NotSupportedError, match="executemany"):
|
|
267
|
+
cur.executemany(obml, [])
|
|
268
|
+
|
|
269
|
+
|
|
270
|
+
def test_cursor_setinputsizes_noop() -> None:
|
|
271
|
+
mock_native = _make_mock_native()
|
|
272
|
+
conn = Connection(mock_native)
|
|
273
|
+
cur = conn.cursor()
|
|
274
|
+
cur.setinputsizes([]) # should not raise
|
|
275
|
+
|
|
276
|
+
|
|
277
|
+
def test_cursor_setoutputsize_noop() -> None:
|
|
278
|
+
mock_native = _make_mock_native()
|
|
279
|
+
conn = Connection(mock_native)
|
|
280
|
+
cur = conn.cursor()
|
|
281
|
+
cur.setoutputsize(1000) # should not raise
|
|
282
|
+
|
|
283
|
+
|
|
284
|
+
def test_cursor_lastrowid_is_none() -> None:
|
|
285
|
+
mock_native = _make_mock_native()
|
|
286
|
+
conn = Connection(mock_native)
|
|
287
|
+
cur = conn.cursor()
|
|
288
|
+
assert cur.lastrowid is None
|
|
289
|
+
|
|
290
|
+
|
|
291
|
+
def test_cursor_rowcount() -> None:
|
|
292
|
+
mock_native = _make_mock_native()
|
|
293
|
+
mock_native.cursor().rowcount = 5
|
|
294
|
+
conn = Connection(mock_native)
|
|
295
|
+
cur = conn.cursor()
|
|
296
|
+
assert cur.rowcount == 5
|
|
297
|
+
|
|
298
|
+
|
|
299
|
+
# ---------------------------------------------------------------------------
|
|
300
|
+
# Cursor — OBML queries (mocked REST API)
|
|
301
|
+
# ---------------------------------------------------------------------------
|
|
302
|
+
|
|
303
|
+
|
|
304
|
+
def test_obml_compile_and_execute() -> None:
|
|
305
|
+
"""OBML query is compiled via REST API then executed on Postgres."""
|
|
306
|
+
mock_native = _make_mock_native()
|
|
307
|
+
compiled_sql = "SELECT region, SUM(amount) AS revenue FROM orders GROUP BY region"
|
|
308
|
+
mock_native.cursor().description = [
|
|
309
|
+
PgColumn("region", 1043, None, None, None, None, None),
|
|
310
|
+
PgColumn("revenue", 701, None, None, None, None, None),
|
|
311
|
+
]
|
|
312
|
+
mock_native.cursor().fetchall.return_value = [
|
|
313
|
+
("EMEA", 300.0), ("APAC", 150.0), ("AMER", 550.0),
|
|
314
|
+
]
|
|
315
|
+
conn = Connection(mock_native)
|
|
316
|
+
with patch("httpx.post", return_value=_mock_api_response(compiled_sql)):
|
|
317
|
+
with conn.cursor() as cur:
|
|
318
|
+
cur.execute("select:\n dimensions:\n - Region\n measures:\n - Revenue\n")
|
|
319
|
+
rows = cur.fetchall()
|
|
320
|
+
assert len(rows) == 3
|
|
321
|
+
# Verify the compiled SQL was passed to native cursor
|
|
322
|
+
mock_native.cursor().execute.assert_called_once_with(compiled_sql)
|
|
323
|
+
|
|
324
|
+
|
|
325
|
+
def test_obml_rest_dialect_is_postgres() -> None:
|
|
326
|
+
"""REST API is called with dialect=postgres."""
|
|
327
|
+
mock_native = _make_mock_native()
|
|
328
|
+
conn = Connection(mock_native)
|
|
329
|
+
compiled_sql = "SELECT 1"
|
|
330
|
+
with patch("httpx.post", return_value=_mock_api_response(compiled_sql)) as mock_post:
|
|
331
|
+
with conn.cursor() as cur:
|
|
332
|
+
cur.execute("select:\n measures:\n - Revenue\n")
|
|
333
|
+
url = mock_post.call_args.args[0]
|
|
334
|
+
assert "/v1/query/sql" in url
|
|
335
|
+
assert mock_post.call_args.kwargs["params"] == {"dialect": "postgres"}
|
|
336
|
+
|
|
337
|
+
|
|
338
|
+
def test_obml_custom_api_url() -> None:
|
|
339
|
+
"""Custom ob_api_url is forwarded to the REST call."""
|
|
340
|
+
mock_native = _make_mock_native()
|
|
341
|
+
conn = Connection(mock_native, ob_api_url="http://my-api:9000")
|
|
342
|
+
compiled_sql = "SELECT 1"
|
|
343
|
+
with patch("httpx.post", return_value=_mock_api_response(compiled_sql)) as mock_post:
|
|
344
|
+
with conn.cursor() as cur:
|
|
345
|
+
cur.execute("select:\n measures:\n - Revenue\n")
|
|
346
|
+
url = mock_post.call_args.args[0]
|
|
347
|
+
assert url == "http://my-api:9000/v1/query/sql"
|
|
348
|
+
|
|
349
|
+
|
|
350
|
+
def test_plain_sql_passthrough() -> None:
|
|
351
|
+
"""Plain SQL is passed through without OBML compilation — no REST call."""
|
|
352
|
+
mock_native = _make_mock_native()
|
|
353
|
+
conn = Connection(mock_native)
|
|
354
|
+
with patch("httpx.post") as mock_post:
|
|
355
|
+
with conn.cursor() as cur:
|
|
356
|
+
cur.execute("SELECT COUNT(*) FROM orders")
|
|
357
|
+
mock_post.assert_not_called()
|
|
358
|
+
mock_native.cursor().execute.assert_called_once_with("SELECT COUNT(*) FROM orders")
|
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
"""Unit tests for OBML YAML detection — no external services required."""
|
|
2
|
+
import pytest
|
|
3
|
+
from ob_postgres.compiler import is_obml, parse_obml
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
@pytest.mark.parametrize("query,expected", [
|
|
7
|
+
("select:\n dimensions:\n - Region\n measures:\n - Revenue\n", True),
|
|
8
|
+
("select:\n measures:\n - Revenue\n", True),
|
|
9
|
+
("select:\n dimensions:\n - Country\n", True),
|
|
10
|
+
("SELECT:\n dimensions:\n - X\n", True),
|
|
11
|
+
(" select:\n dimensions:\n - X\n", True),
|
|
12
|
+
("SELECT * FROM orders", False),
|
|
13
|
+
("SELECT id\nFROM customers WHERE id = 1", False),
|
|
14
|
+
("", False),
|
|
15
|
+
("model:\n name: test\n", False),
|
|
16
|
+
("select:\n filters:\n - x = 1\n", False),
|
|
17
|
+
("select:\n - [unclosed", False),
|
|
18
|
+
("select:\n - dimension1\n", False),
|
|
19
|
+
])
|
|
20
|
+
def test_obml_detection(query, expected):
|
|
21
|
+
assert is_obml(query) == expected
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def test_parse_obml_returns_dict():
|
|
25
|
+
q = "select:\n dimensions:\n - Region\n measures:\n - Revenue\nlimit: 100\n"
|
|
26
|
+
result = parse_obml(q)
|
|
27
|
+
assert result["select"]["dimensions"] == ["Region"]
|
|
28
|
+
assert result["select"]["measures"] == ["Revenue"]
|
|
29
|
+
assert result["limit"] == 100
|