tapestry-orm 0.0.1__tar.gz → 0.0.2__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.
- {tapestry_orm-0.0.1 → tapestry_orm-0.0.2}/PKG-INFO +15 -5
- {tapestry_orm-0.0.1 → tapestry_orm-0.0.2}/README.md +2 -2
- {tapestry_orm-0.0.1 → tapestry_orm-0.0.2}/pyproject.toml +19 -17
- tapestry_orm-0.0.2/src/tapestry/__init__.py +11 -0
- {tapestry_orm-0.0.1 → tapestry_orm-0.0.2}/src/tapestry/base.py +23 -19
- {tapestry_orm-0.0.1 → tapestry_orm-0.0.2}/src/tapestry/edge.py +3 -3
- tapestry_orm-0.0.2/src/tapestry/engine.py +398 -0
- {tapestry_orm-0.0.1 → tapestry_orm-0.0.2}/src/tapestry/field.py +24 -27
- tapestry_orm-0.0.2/src/tapestry/field.pyi +100 -0
- {tapestry_orm-0.0.1 → tapestry_orm-0.0.2}/src/tapestry/query.py +10 -16
- {tapestry_orm-0.0.1 → tapestry_orm-0.0.2}/src/tapestry/table.py +9 -11
- {tapestry_orm-0.0.1 → tapestry_orm-0.0.2}/src/tapestry/tokenizer.py +2 -4
- {tapestry_orm-0.0.1 → tapestry_orm-0.0.2}/src/tapestry/utils.py +8 -8
- tapestry_orm-0.0.1/src/tapestry/__init__.py +0 -11
- tapestry_orm-0.0.1/src/tapestry/base.pyi +0 -54
- tapestry_orm-0.0.1/src/tapestry/edge.pyi +0 -45
- tapestry_orm-0.0.1/src/tapestry/engine.py +0 -301
- tapestry_orm-0.0.1/src/tapestry/field.pyi +0 -113
- tapestry_orm-0.0.1/src/tapestry/node.pyi +0 -28
- tapestry_orm-0.0.1/src/tapestry/query.pyi +0 -33
- {tapestry_orm-0.0.1 → tapestry_orm-0.0.2}/src/tapestry/node.py +0 -0
- {tapestry_orm-0.0.1 → tapestry_orm-0.0.2}/src/tapestry/py.typed +0 -0
|
@@ -1,16 +1,26 @@
|
|
|
1
|
-
Metadata-Version: 2.
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
2
|
Name: tapestry-orm
|
|
3
|
-
Version: 0.0.
|
|
3
|
+
Version: 0.0.2
|
|
4
4
|
Summary: A prototype ORM for SurrealDB
|
|
5
5
|
Author-email: timothee@obrecht.xyz
|
|
6
|
-
License: MIT
|
|
6
|
+
License-Expression: MIT
|
|
7
7
|
Requires-Dist: more-itertools>=10.7.0
|
|
8
8
|
Requires-Dist: pydantic>=2.11.7
|
|
9
9
|
Requires-Dist: surrealdb>=1.0.6
|
|
10
|
+
Requires-Dist: sphinx>=7.2.0 ; extra == 'docs'
|
|
11
|
+
Requires-Dist: sphinx-rtd-theme>=2.0.0 ; extra == 'docs'
|
|
12
|
+
Requires-Dist: sphinx-autodoc-typehints>=2.0.0 ; extra == 'docs'
|
|
13
|
+
Requires-Dist: sphinx-copybutton>=0.5.2 ; extra == 'docs'
|
|
14
|
+
Requires-Dist: myst-parser>=2.0.0 ; extra == 'docs'
|
|
15
|
+
Requires-Dist: sphinx-autobuild>=2021.3.14 ; extra == 'docs'
|
|
16
|
+
Requires-Dist: sphinx-autoapi>=3.0.0 ; extra == 'docs'
|
|
10
17
|
Requires-Python: >=3.12
|
|
18
|
+
Project-URL: Documentation, https://readthedocs.org
|
|
19
|
+
Project-URL: Repository, https://gitlab.com/timete_OS/tapestry-orm
|
|
20
|
+
Provides-Extra: docs
|
|
11
21
|
Description-Content-Type: text/markdown
|
|
12
22
|
|
|
13
|
-
#
|
|
23
|
+
# Tapestry
|
|
14
24
|
|
|
15
25
|
**A modern, type-safe Python ORM for SurrealDB**
|
|
16
26
|
|
|
@@ -341,7 +351,7 @@ To test the GitLab CI pipeline locally before pushing, use [`gitlab-ci-local`](h
|
|
|
341
351
|
|
|
342
352
|
```bash
|
|
343
353
|
# Run the test job locally (configuration is already set up)
|
|
344
|
-
gitlab-ci-local
|
|
354
|
+
gitlab-ci-local
|
|
345
355
|
```
|
|
346
356
|
|
|
347
357
|
The project includes two configuration files that make this work:
|
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
#
|
|
1
|
+
# Tapestry
|
|
2
2
|
|
|
3
3
|
**A modern, type-safe Python ORM for SurrealDB**
|
|
4
4
|
|
|
@@ -329,7 +329,7 @@ To test the GitLab CI pipeline locally before pushing, use [`gitlab-ci-local`](h
|
|
|
329
329
|
|
|
330
330
|
```bash
|
|
331
331
|
# Run the test job locally (configuration is already set up)
|
|
332
|
-
gitlab-ci-local
|
|
332
|
+
gitlab-ci-local
|
|
333
333
|
```
|
|
334
334
|
|
|
335
335
|
The project includes two configuration files that make this work:
|
|
@@ -1,21 +1,38 @@
|
|
|
1
1
|
[project]
|
|
2
2
|
name = "tapestry-orm"
|
|
3
|
-
version = "0.0.
|
|
3
|
+
version = "0.0.2"
|
|
4
4
|
description = "A prototype ORM for SurrealDB"
|
|
5
5
|
readme = "README.md"
|
|
6
|
-
license =
|
|
6
|
+
license = "MIT"
|
|
7
|
+
license-files = ["licence.txt"]
|
|
7
8
|
authors = [
|
|
8
9
|
{ email = "timothee@obrecht.xyz" }
|
|
9
10
|
]
|
|
11
|
+
project_page = "https://gitlab.com/timete_OS/tapestry-orm"
|
|
10
12
|
requires-python = ">=3.12"
|
|
11
13
|
dependencies = [
|
|
12
14
|
"more-itertools>=10.7.0",
|
|
13
15
|
"pydantic>=2.11.7",
|
|
14
16
|
"surrealdb>=1.0.6",
|
|
15
17
|
]
|
|
18
|
+
[project.urls]
|
|
19
|
+
Repository = "https://gitlab.com/timete_OS/tapestry-orm"
|
|
20
|
+
Documentation = "https://readthedocs.org"
|
|
21
|
+
|
|
22
|
+
[project.optional-dependencies]
|
|
23
|
+
docs = [
|
|
24
|
+
"sphinx>=7.2.0",
|
|
25
|
+
"sphinx-rtd-theme>=2.0.0",
|
|
26
|
+
"sphinx-autodoc-typehints>=2.0.0",
|
|
27
|
+
"sphinx-copybutton>=0.5.2",
|
|
28
|
+
"myst-parser>=2.0.0",
|
|
29
|
+
"sphinx-autobuild>=2021.3.14",
|
|
30
|
+
"sphinx-autoapi>=3.0.0",
|
|
31
|
+
]
|
|
16
32
|
|
|
17
33
|
[dependency-groups]
|
|
18
34
|
dev = [
|
|
35
|
+
"basedpyright>=1.32.1",
|
|
19
36
|
"pytest>=8.4.1",
|
|
20
37
|
"pytest-asyncio>=1.1.0",
|
|
21
38
|
"testcontainers>=4.12.0",
|
|
@@ -25,13 +42,7 @@ dev = [
|
|
|
25
42
|
misc = [
|
|
26
43
|
"pandas>=2.3.1",
|
|
27
44
|
"ipykernel>=6.30.1",
|
|
28
|
-
"sphinx>=7.2.0",
|
|
29
|
-
"sphinx-rtd-theme>=2.0.0",
|
|
30
|
-
"sphinx-autodoc-typehints>=2.0.0",
|
|
31
|
-
"sphinx-copybutton>=0.5.2",
|
|
32
|
-
"myst-parser>=2.0.0",
|
|
33
45
|
"mypy>=1.18.2",
|
|
34
|
-
"pyright>=1.1.406",
|
|
35
46
|
]
|
|
36
47
|
|
|
37
48
|
[build-system]
|
|
@@ -41,16 +52,7 @@ build-backend = "uv_build"
|
|
|
41
52
|
[tool.uv.build-backend]
|
|
42
53
|
module-name = "tapestry"
|
|
43
54
|
|
|
44
|
-
[tool.uv.sources]
|
|
45
|
-
# territories = { path = "../territories", editable = true }
|
|
46
|
-
|
|
47
55
|
[tool.pytest.ini_options]
|
|
48
56
|
filterwarnings = [
|
|
49
57
|
"ignore::DeprecationWarning:testcontainers.*",
|
|
50
58
|
]
|
|
51
|
-
|
|
52
|
-
[[tool.uv.index]]
|
|
53
|
-
name = "datapolitics"
|
|
54
|
-
url = "http://pypi.hedwige.fr:8080"
|
|
55
|
-
publish-url = "http://pypi.hedwige.fr:8080"
|
|
56
|
-
allow-insecure-host = ["http://pypi.hedwige.fr:8080"]
|
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
__all__ = [
|
|
2
|
+
"Base", "Text", "create_engine", "create_engine_context", "Engine", "Reference", "Node", "Edge", "Q"
|
|
3
|
+
]
|
|
4
|
+
|
|
5
|
+
from .query import Q
|
|
6
|
+
from .base import Base
|
|
7
|
+
from .node import Node
|
|
8
|
+
from .edge import Edge
|
|
9
|
+
from .tokenizer import Text
|
|
10
|
+
from .table import Reference
|
|
11
|
+
from .engine import create_engine, create_engine_context, Engine
|
|
@@ -1,19 +1,21 @@
|
|
|
1
1
|
from __future__ import annotations
|
|
2
2
|
|
|
3
|
+
from collections.abc import Mapping
|
|
3
4
|
import inspect
|
|
5
|
+
import types
|
|
4
6
|
|
|
5
7
|
from enum import Enum
|
|
6
8
|
from copy import deepcopy
|
|
7
9
|
from itertools import chain
|
|
8
10
|
from queue import LifoQueue
|
|
11
|
+
from surrealdb import RecordID
|
|
12
|
+
from typing import Any, ClassVar, Unpack, get_origin, get_args
|
|
9
13
|
from pydantic_core.core_schema import ValidationInfo
|
|
10
14
|
from pydantic.fields import FieldInfo, ComputedFieldInfo
|
|
11
|
-
from
|
|
12
|
-
from typing import Any, ClassVar, Optional, Self, Any, Unpack, Union, Type, get_args
|
|
13
|
-
from pydantic import BaseModel, field_validator, model_serializer, field_serializer, Field, ConfigDict, model_validator
|
|
15
|
+
from pydantic import BaseModel, field_validator, model_serializer, Field, ConfigDict
|
|
14
16
|
|
|
15
|
-
from .table import Table
|
|
16
|
-
from .utils import convert_types, replace_type
|
|
17
|
+
from .table import Table
|
|
18
|
+
from .utils import convert_types, replace_type
|
|
17
19
|
from .field import Field as CustomField, NestedFieldDescriptor, ComputedFieldDescriptor
|
|
18
20
|
|
|
19
21
|
|
|
@@ -57,20 +59,21 @@ class Base(BaseModel):
|
|
|
57
59
|
"""
|
|
58
60
|
id: RecordID | None = Field(exclude=False, default = None)
|
|
59
61
|
|
|
60
|
-
model_config =
|
|
61
|
-
|
|
62
|
-
|
|
63
|
-
|
|
64
|
-
|
|
62
|
+
model_config: ClassVar[ConfigDict] = ConfigDict(
|
|
63
|
+
arbitrary_types_allowed=True,
|
|
64
|
+
use_enum_values=False, # would be great to have a 'use_enum_names'
|
|
65
|
+
validate_assignment=True
|
|
66
|
+
)
|
|
65
67
|
_registry: ClassVar[tuple[Table, ...]] = ()
|
|
66
68
|
_tokenizers: ClassVar[set[str]] = set()
|
|
67
69
|
_to_create: ClassVar[LifoQueue[Base]]
|
|
68
|
-
|
|
70
|
+
_discarded: ClassVar[bool] = False
|
|
71
|
+
child_classes: ClassVar[dict[str, type[Base]]] = {}
|
|
69
72
|
|
|
70
73
|
|
|
71
74
|
@classmethod
|
|
72
75
|
@field_validator('id')
|
|
73
|
-
def validate_record_id(cls, v: Any) ->
|
|
76
|
+
def validate_record_id(cls, v: Any) -> RecordID | None:
|
|
74
77
|
"""
|
|
75
78
|
Validate and convert record ID values.
|
|
76
79
|
|
|
@@ -168,7 +171,7 @@ class Base(BaseModel):
|
|
|
168
171
|
# Create unique queue for each subclass
|
|
169
172
|
# this queue will be used to chain several object creations
|
|
170
173
|
cls._to_create = LifoQueue()
|
|
171
|
-
cls._discarded = discarded
|
|
174
|
+
cls._discarded = bool(discarded)
|
|
172
175
|
return super().__init_subclass__(**kwargs)
|
|
173
176
|
|
|
174
177
|
|
|
@@ -180,8 +183,9 @@ class Base(BaseModel):
|
|
|
180
183
|
if cls._discarded:
|
|
181
184
|
return
|
|
182
185
|
|
|
183
|
-
|
|
184
|
-
|
|
186
|
+
model_fields_copy: dict[str, FieldInfo | ComputedFieldInfo] = dict(deepcopy(cls.model_fields))
|
|
187
|
+
model_fields_copy.update(cls.model_computed_fields) # type: ignore[arg-type]
|
|
188
|
+
model_fields: Mapping[str, FieldInfo | ComputedFieldInfo] = model_fields_copy
|
|
185
189
|
fields: list[CustomField[Any]] = []
|
|
186
190
|
for field_name, field_info in model_fields.items():
|
|
187
191
|
field_type = field_info.annotation if isinstance(field_info, FieldInfo) else field_info.return_type
|
|
@@ -198,9 +202,9 @@ class Base(BaseModel):
|
|
|
198
202
|
if field_type and hasattr(field_type, '__mro__') and Base in field_type.__mro__:
|
|
199
203
|
nested_class = field_type
|
|
200
204
|
# Check if it's Optional[Base subclass] (Union with None)
|
|
201
|
-
elif hasattr(field_type, '
|
|
205
|
+
elif get_origin(field_type) is types.UnionType or (get_origin(field_type) is not None and hasattr(get_origin(field_type), '__name__') and get_origin(field_type).__name__ == 'UnionType'):
|
|
202
206
|
# Get the args of the Union
|
|
203
|
-
args =
|
|
207
|
+
args = get_args(field_type)
|
|
204
208
|
# Filter out None and check if any remaining type is a Base subclass
|
|
205
209
|
for arg in args:
|
|
206
210
|
if arg is not type(None) and hasattr(arg, '__mro__') and Base in arg.__mro__:
|
|
@@ -219,7 +223,7 @@ class Base(BaseModel):
|
|
|
219
223
|
|
|
220
224
|
# replace Base instance with Union[annotation, RecordID]
|
|
221
225
|
if field_name in cls.model_fields:
|
|
222
|
-
cls.model_fields[field_name].annotation = replace_type(cls.model_fields[field_name].annotation, Base,
|
|
226
|
+
cls.model_fields[field_name].annotation = replace_type(cls.model_fields[field_name].annotation, Base, Base | RecordID)
|
|
223
227
|
|
|
224
228
|
|
|
225
229
|
table = Table(name=cls.__name__.lower(), base_class=Base, model_class=cls, fields=tuple(fields), relation=relation)
|
|
@@ -249,7 +253,7 @@ class Base(BaseModel):
|
|
|
249
253
|
|
|
250
254
|
|
|
251
255
|
@classmethod
|
|
252
|
-
def add_table(cls, table: Table, child_class:
|
|
256
|
+
def add_table(cls, table: Table, child_class: type[Base]):
|
|
253
257
|
"""
|
|
254
258
|
Register a new table in the system registry.
|
|
255
259
|
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
from abc import ABC
|
|
2
2
|
from copy import deepcopy
|
|
3
|
-
from typing import Self, Iterable, get_type_hints, Any, ClassVar,
|
|
3
|
+
from typing import Self, Iterable, get_type_hints, Any, ClassVar, TYPE_CHECKING
|
|
4
4
|
from pydantic.fields import FieldInfo, ComputedFieldInfo
|
|
5
5
|
from surrealdb import AsyncWsSurrealConnection, AsyncHttpSurrealConnection, RecordID
|
|
6
6
|
|
|
@@ -113,8 +113,8 @@ class Edge(Base, ABC, discarded=True):
|
|
|
113
113
|
super().__pydantic_init_subclass__(relation=relation, **kwargs)
|
|
114
114
|
# need to add those back for validation at instanciation
|
|
115
115
|
# otherwise, the field is skipped
|
|
116
|
-
in_.annotation = replace_type(in_.annotation, Base,
|
|
117
|
-
out_.annotation = replace_type(out_.annotation, Base,
|
|
116
|
+
in_.annotation = replace_type(in_.annotation, Base, Base | RecordID)
|
|
117
|
+
out_.annotation = replace_type(out_.annotation, Base, Base | RecordID)
|
|
118
118
|
cls.model_fields["in_"] = in_
|
|
119
119
|
cls.model_fields["out_"] = out_
|
|
120
120
|
cls.model_rebuild(force=True)
|
|
@@ -0,0 +1,398 @@
|
|
|
1
|
+
import asyncio
|
|
2
|
+
import logging
|
|
3
|
+
import contextlib
|
|
4
|
+
from typing import Optional
|
|
5
|
+
|
|
6
|
+
from collections.abc import AsyncIterator
|
|
7
|
+
from surrealdb import AsyncSurreal, AsyncWsSurrealConnection, AsyncHttpSurrealConnection
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
logger = logging.getLogger(__name__)
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class Engine:
|
|
14
|
+
"""
|
|
15
|
+
SurrealDB connection pool engine, similar to SQLAlchemy's engine.
|
|
16
|
+
|
|
17
|
+
This class manages a pool of authenticated SurrealDB connections that can be
|
|
18
|
+
reused across async operations. It provides a clean API for connection management
|
|
19
|
+
with context managers.
|
|
20
|
+
|
|
21
|
+
The Engine should be created with `create_engine()` and manages its own lifecycle.
|
|
22
|
+
|
|
23
|
+
Attributes:
|
|
24
|
+
url: SurrealDB connection URL
|
|
25
|
+
auth_payload: Authentication credentials dictionary
|
|
26
|
+
pool_size: Number of connections to maintain in the pool
|
|
27
|
+
connect_timeout: Timeout for connection attempts in seconds
|
|
28
|
+
|
|
29
|
+
Example:
|
|
30
|
+
>>> # Create and start engine
|
|
31
|
+
>>> engine = create_engine(
|
|
32
|
+
... url="ws://localhost:8000/rpc",
|
|
33
|
+
... auth_payload={"username": "root", "password": "root"},
|
|
34
|
+
... pool_size=5
|
|
35
|
+
... )
|
|
36
|
+
>>> await engine.start()
|
|
37
|
+
>>>
|
|
38
|
+
>>> # Use connections via context manager
|
|
39
|
+
>>> async with engine.connect() as db:
|
|
40
|
+
... await db.use("mydb", "myns")
|
|
41
|
+
... results = await db.select("person")
|
|
42
|
+
>>>
|
|
43
|
+
>>> # Close when done
|
|
44
|
+
>>> await engine.close()
|
|
45
|
+
|
|
46
|
+
Or use with lifecycle context manager:
|
|
47
|
+
>>> async with create_engine_context(url, auth_payload) as engine:
|
|
48
|
+
... async with engine.connect() as db:
|
|
49
|
+
... # use db
|
|
50
|
+
... pass
|
|
51
|
+
"""
|
|
52
|
+
|
|
53
|
+
def __init__(
|
|
54
|
+
self,
|
|
55
|
+
url: str,
|
|
56
|
+
auth_payload: dict,
|
|
57
|
+
pool_size: int = 4,
|
|
58
|
+
connect_timeout: float = 5.0,
|
|
59
|
+
):
|
|
60
|
+
"""
|
|
61
|
+
Initialize the engine.
|
|
62
|
+
|
|
63
|
+
Note: The engine is not started automatically. Call await engine.start()
|
|
64
|
+
or use create_engine_context() for automatic lifecycle management.
|
|
65
|
+
|
|
66
|
+
Args:
|
|
67
|
+
url: SurrealDB connection URL (ws:// or http://)
|
|
68
|
+
auth_payload: Dictionary with authentication credentials
|
|
69
|
+
pool_size: Number of connections to maintain (default: 4)
|
|
70
|
+
connect_timeout: Connection timeout in seconds (default: 5.0)
|
|
71
|
+
"""
|
|
72
|
+
self.url = url
|
|
73
|
+
self.auth_payload = auth_payload
|
|
74
|
+
self.pool_size = pool_size
|
|
75
|
+
self.connect_timeout = connect_timeout
|
|
76
|
+
self._queue: Optional[asyncio.Queue] = None
|
|
77
|
+
self._clients: list = []
|
|
78
|
+
self._closed = False
|
|
79
|
+
self._started = False
|
|
80
|
+
|
|
81
|
+
async def _make_client(self) -> AsyncSurreal:
|
|
82
|
+
"""
|
|
83
|
+
Create and authenticate a single SurrealDB client.
|
|
84
|
+
|
|
85
|
+
Returns:
|
|
86
|
+
An authenticated AsyncSurreal client ready for use
|
|
87
|
+
|
|
88
|
+
Raises:
|
|
89
|
+
Exception: If connection or authentication fails
|
|
90
|
+
"""
|
|
91
|
+
client = AsyncSurreal(self.url)
|
|
92
|
+
await client.signin(self.auth_payload)
|
|
93
|
+
return client
|
|
94
|
+
|
|
95
|
+
async def start(self) -> "Engine":
|
|
96
|
+
"""
|
|
97
|
+
Initialize the connection pool with authenticated clients.
|
|
98
|
+
|
|
99
|
+
Creates the specified number of clients, authenticates them,
|
|
100
|
+
and adds them to the pool queue for reuse.
|
|
101
|
+
|
|
102
|
+
Returns:
|
|
103
|
+
self: Returns the engine instance for chaining
|
|
104
|
+
|
|
105
|
+
Raises:
|
|
106
|
+
RuntimeError: If the engine is already started or closed
|
|
107
|
+
Exception: If unable to create the required number of clients
|
|
108
|
+
"""
|
|
109
|
+
if self._started:
|
|
110
|
+
raise RuntimeError("Engine already started")
|
|
111
|
+
if self._closed:
|
|
112
|
+
raise RuntimeError("Engine is closed")
|
|
113
|
+
|
|
114
|
+
self._queue = asyncio.Queue()
|
|
115
|
+
|
|
116
|
+
for i in range(self.pool_size):
|
|
117
|
+
try:
|
|
118
|
+
client = await self._make_client()
|
|
119
|
+
except Exception as e:
|
|
120
|
+
logger.exception("Failed to create surreal client #%d: %s", i, e)
|
|
121
|
+
# Retry once with backoff before giving up
|
|
122
|
+
await asyncio.sleep(0.5)
|
|
123
|
+
client = await self._make_client()
|
|
124
|
+
|
|
125
|
+
self._clients.append(client)
|
|
126
|
+
self._queue.put_nowait(client)
|
|
127
|
+
|
|
128
|
+
self._started = True
|
|
129
|
+
logger.info("SurrealDB engine started with %d connections", len(self._clients))
|
|
130
|
+
return self
|
|
131
|
+
|
|
132
|
+
async def close(self) -> None:
|
|
133
|
+
"""
|
|
134
|
+
Close all connections and shut down the engine.
|
|
135
|
+
|
|
136
|
+
Drains the queue and closes all client connections gracefully.
|
|
137
|
+
This should be called when the application shuts down.
|
|
138
|
+
"""
|
|
139
|
+
if self._closed:
|
|
140
|
+
return
|
|
141
|
+
|
|
142
|
+
self._closed = True
|
|
143
|
+
|
|
144
|
+
# Drain queue
|
|
145
|
+
if self._queue is not None:
|
|
146
|
+
while not self._queue.empty():
|
|
147
|
+
try:
|
|
148
|
+
_ = self._queue.get_nowait()
|
|
149
|
+
self._queue.task_done()
|
|
150
|
+
except asyncio.QueueEmpty:
|
|
151
|
+
break
|
|
152
|
+
|
|
153
|
+
# Close all clients
|
|
154
|
+
for client in self._clients:
|
|
155
|
+
try:
|
|
156
|
+
await client.close()
|
|
157
|
+
except Exception:
|
|
158
|
+
logger.exception("Error closing surreal client")
|
|
159
|
+
|
|
160
|
+
self._clients.clear()
|
|
161
|
+
self._started = False
|
|
162
|
+
logger.info("SurrealDB engine closed")
|
|
163
|
+
|
|
164
|
+
@contextlib.asynccontextmanager
|
|
165
|
+
async def connect(self) -> AsyncIterator[AsyncWsSurrealConnection | AsyncHttpSurrealConnection]:
|
|
166
|
+
"""
|
|
167
|
+
Acquire a connection from the pool as a context manager.
|
|
168
|
+
|
|
169
|
+
This is the primary way to get a database connection from the engine.
|
|
170
|
+
The connection is automatically returned to the pool when the context exits.
|
|
171
|
+
|
|
172
|
+
Yields:
|
|
173
|
+
AsyncWsSurrealConnection | AsyncHttpSurrealConnection: An authenticated client
|
|
174
|
+
|
|
175
|
+
Example:
|
|
176
|
+
>>> async with engine.connect() as db:
|
|
177
|
+
... await db.use("mydb", "myns")
|
|
178
|
+
... results = await db.select("person")
|
|
179
|
+
|
|
180
|
+
Raises:
|
|
181
|
+
RuntimeError: If the engine is not started or is closed
|
|
182
|
+
Exception: If the acquired client encounters a fatal error
|
|
183
|
+
"""
|
|
184
|
+
if self._closed:
|
|
185
|
+
raise RuntimeError("Engine is closed")
|
|
186
|
+
if not self._started:
|
|
187
|
+
raise RuntimeError("Engine not started. Call await engine.start() first")
|
|
188
|
+
|
|
189
|
+
client = await self._queue.get()
|
|
190
|
+
try:
|
|
191
|
+
yield client
|
|
192
|
+
except Exception as exc:
|
|
193
|
+
# If the client had a fatal error, replace it
|
|
194
|
+
if self._is_fatal_client_error(exc):
|
|
195
|
+
logger.warning("Client had fatal error, replacing: %s", exc)
|
|
196
|
+
await self._replace_client(client)
|
|
197
|
+
raise
|
|
198
|
+
finally:
|
|
199
|
+
# Return to pool if not closed
|
|
200
|
+
if not self._closed and self._queue is not None:
|
|
201
|
+
self._queue.put_nowait(client)
|
|
202
|
+
|
|
203
|
+
# Alias for SQLAlchemy-like API
|
|
204
|
+
acquire = connect
|
|
205
|
+
|
|
206
|
+
async def execute(self, query: str, *args, retries: int = 1, backoff: float = 0.2, **kwargs):
|
|
207
|
+
"""
|
|
208
|
+
Execute a query using a pooled connection with automatic retry.
|
|
209
|
+
|
|
210
|
+
Acquires a client from the pool, executes the query, and handles
|
|
211
|
+
retries for transient failures.
|
|
212
|
+
|
|
213
|
+
Args:
|
|
214
|
+
query: The query string to execute
|
|
215
|
+
*args: Arguments to pass to the query method
|
|
216
|
+
retries: Number of retry attempts for transient errors (default: 1)
|
|
217
|
+
backoff: Base backoff time in seconds between retries (default: 0.2)
|
|
218
|
+
**kwargs: Keyword arguments to pass to the query method
|
|
219
|
+
|
|
220
|
+
Returns:
|
|
221
|
+
Query results from SurrealDB
|
|
222
|
+
|
|
223
|
+
Raises:
|
|
224
|
+
RuntimeError: If the engine is not started
|
|
225
|
+
Exception: If the query fails after all retry attempts
|
|
226
|
+
"""
|
|
227
|
+
if self._closed:
|
|
228
|
+
raise RuntimeError("Engine is closed")
|
|
229
|
+
if not self._started:
|
|
230
|
+
raise RuntimeError("Engine not started. Call await engine.start() first")
|
|
231
|
+
|
|
232
|
+
for attempt in range(retries + 1):
|
|
233
|
+
async with self.connect() as client:
|
|
234
|
+
try:
|
|
235
|
+
return await client.query(query, *args, **kwargs)
|
|
236
|
+
except Exception as exc:
|
|
237
|
+
logger.exception("Query failed on attempt %d: %s", attempt, exc)
|
|
238
|
+
if attempt < retries and self._is_retryable_error(exc):
|
|
239
|
+
await asyncio.sleep(backoff * (2 ** attempt))
|
|
240
|
+
continue
|
|
241
|
+
raise
|
|
242
|
+
|
|
243
|
+
async def _replace_client(self, dead_client: AsyncSurreal) -> None:
|
|
244
|
+
"""
|
|
245
|
+
Replace a failed client with a new authenticated connection.
|
|
246
|
+
|
|
247
|
+
Closes the dead client, creates a new one, and adds it back to the pool
|
|
248
|
+
to maintain the pool size.
|
|
249
|
+
|
|
250
|
+
Args:
|
|
251
|
+
dead_client: The client connection that failed
|
|
252
|
+
"""
|
|
253
|
+
try:
|
|
254
|
+
await dead_client.close()
|
|
255
|
+
except Exception:
|
|
256
|
+
pass
|
|
257
|
+
|
|
258
|
+
try:
|
|
259
|
+
new_client = await self._make_client()
|
|
260
|
+
except Exception as e:
|
|
261
|
+
logger.exception("Failed to create replacement client: %s", e)
|
|
262
|
+
# If we can't create a replacement, put back the old one to avoid pool starvation
|
|
263
|
+
if self._queue is not None and not self._closed:
|
|
264
|
+
self._queue.put_nowait(dead_client)
|
|
265
|
+
return
|
|
266
|
+
|
|
267
|
+
# Replace in internal list
|
|
268
|
+
try:
|
|
269
|
+
idx = self._clients.index(dead_client)
|
|
270
|
+
self._clients[idx] = new_client
|
|
271
|
+
except ValueError:
|
|
272
|
+
self._clients.append(new_client)
|
|
273
|
+
|
|
274
|
+
if self._queue is not None and not self._closed:
|
|
275
|
+
self._queue.put_nowait(new_client)
|
|
276
|
+
|
|
277
|
+
def _is_retryable_error(self, exc: Exception) -> bool:
|
|
278
|
+
"""
|
|
279
|
+
Determine if an exception is retryable.
|
|
280
|
+
|
|
281
|
+
Args:
|
|
282
|
+
exc: The exception to check
|
|
283
|
+
|
|
284
|
+
Returns:
|
|
285
|
+
bool: True if the error is transient and worth retrying
|
|
286
|
+
"""
|
|
287
|
+
return isinstance(exc, (asyncio.TimeoutError, ConnectionError))
|
|
288
|
+
|
|
289
|
+
def _is_fatal_client_error(self, exc: Exception) -> bool:
|
|
290
|
+
"""
|
|
291
|
+
Determine if an exception indicates a broken client connection.
|
|
292
|
+
|
|
293
|
+
Args:
|
|
294
|
+
exc: The exception to check
|
|
295
|
+
|
|
296
|
+
Returns:
|
|
297
|
+
bool: True if the client connection is likely broken and needs replacement
|
|
298
|
+
"""
|
|
299
|
+
return isinstance(exc, (ConnectionError, RuntimeError))
|
|
300
|
+
|
|
301
|
+
|
|
302
|
+
def create_engine(
|
|
303
|
+
url: str,
|
|
304
|
+
auth_payload: dict,
|
|
305
|
+
pool_size: int = 4,
|
|
306
|
+
connect_timeout: float = 5.0,
|
|
307
|
+
) -> Engine:
|
|
308
|
+
"""
|
|
309
|
+
Create a connection pool engine for SurrealDB.
|
|
310
|
+
|
|
311
|
+
This function creates an Engine instance that manages a pool of authenticated
|
|
312
|
+
SurrealDB connections. The engine must be started with await engine.start()
|
|
313
|
+
before use.
|
|
314
|
+
|
|
315
|
+
Args:
|
|
316
|
+
url: SurrealDB connection URL (e.g., "ws://localhost:8000/rpc")
|
|
317
|
+
auth_payload: Dictionary with authentication credentials
|
|
318
|
+
Example: {"username": "root", "password": "root"}
|
|
319
|
+
pool_size: Number of connections to maintain in the pool (default: 4)
|
|
320
|
+
connect_timeout: Timeout for connection attempts in seconds (default: 5.0)
|
|
321
|
+
|
|
322
|
+
Returns:
|
|
323
|
+
Engine: An engine instance (not yet started)
|
|
324
|
+
|
|
325
|
+
Example:
|
|
326
|
+
>>> # Create and manually manage lifecycle
|
|
327
|
+
>>> engine = create_engine(
|
|
328
|
+
... "ws://localhost:8000/rpc",
|
|
329
|
+
... {"username": "root", "password": "root"},
|
|
330
|
+
... pool_size=5
|
|
331
|
+
... )
|
|
332
|
+
>>> await engine.start()
|
|
333
|
+
>>>
|
|
334
|
+
>>> async with engine.connect() as db:
|
|
335
|
+
... await db.use("mydatabase", "mynamespace")
|
|
336
|
+
... people = await Person.insert(db, [
|
|
337
|
+
... Person(name="Alice"),
|
|
338
|
+
... Person(name="Bob")
|
|
339
|
+
... ])
|
|
340
|
+
>>>
|
|
341
|
+
>>> await engine.close()
|
|
342
|
+
|
|
343
|
+
See Also:
|
|
344
|
+
- create_engine_context(): For automatic lifecycle management
|
|
345
|
+
- Engine.connect(): Get a connection from the pool
|
|
346
|
+
- Engine.execute(): Execute a query with automatic retry
|
|
347
|
+
"""
|
|
348
|
+
return Engine(url, auth_payload, pool_size, connect_timeout)
|
|
349
|
+
|
|
350
|
+
|
|
351
|
+
@contextlib.asynccontextmanager
|
|
352
|
+
async def create_engine_context(
|
|
353
|
+
url: str,
|
|
354
|
+
auth_payload: dict,
|
|
355
|
+
pool_size: int = 4,
|
|
356
|
+
connect_timeout: float = 5.0,
|
|
357
|
+
) -> AsyncIterator[Engine]:
|
|
358
|
+
"""
|
|
359
|
+
Create a connection pool engine with automatic lifecycle management.
|
|
360
|
+
|
|
361
|
+
This context manager creates an engine, starts it, yields it for use,
|
|
362
|
+
and automatically closes it when the context exits. This is the recommended
|
|
363
|
+
way to use the engine as it ensures proper cleanup.
|
|
364
|
+
|
|
365
|
+
Args:
|
|
366
|
+
url: SurrealDB connection URL (e.g., "ws://localhost:8000/rpc")
|
|
367
|
+
auth_payload: Dictionary with authentication credentials
|
|
368
|
+
Example: {"username": "root", "password": "root"}
|
|
369
|
+
pool_size: Number of connections to maintain in the pool (default: 4)
|
|
370
|
+
connect_timeout: Timeout for connection attempts in seconds (default: 5.0)
|
|
371
|
+
|
|
372
|
+
Yields:
|
|
373
|
+
Engine: A started engine instance
|
|
374
|
+
|
|
375
|
+
Example:
|
|
376
|
+
>>> async with create_engine_context(
|
|
377
|
+
... "ws://localhost:8000/rpc",
|
|
378
|
+
... {"username": "root", "password": "root"},
|
|
379
|
+
... pool_size=5
|
|
380
|
+
... ) as engine:
|
|
381
|
+
... async with engine.connect() as db:
|
|
382
|
+
... await db.use("mydatabase", "mynamespace")
|
|
383
|
+
... people = await Person.insert(db, [
|
|
384
|
+
... Person(name="Alice"),
|
|
385
|
+
... Person(name="Bob")
|
|
386
|
+
... ])
|
|
387
|
+
... # Engine is automatically closed here
|
|
388
|
+
|
|
389
|
+
See Also:
|
|
390
|
+
- create_engine(): For manual lifecycle management
|
|
391
|
+
- Engine.connect(): Get a connection from the pool
|
|
392
|
+
"""
|
|
393
|
+
engine = Engine(url, auth_payload, pool_size, connect_timeout)
|
|
394
|
+
await engine.start()
|
|
395
|
+
try:
|
|
396
|
+
yield engine
|
|
397
|
+
finally:
|
|
398
|
+
await engine.close()
|