dbt-redis-adbc 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.
- dbt_redis_adbc-0.1.0/.gitignore +11 -0
- dbt_redis_adbc-0.1.0/PKG-INFO +88 -0
- dbt_redis_adbc-0.1.0/README.md +67 -0
- dbt_redis_adbc-0.1.0/dbt/adapters/redis_adbc/__init__.py +13 -0
- dbt_redis_adbc-0.1.0/dbt/adapters/redis_adbc/__version__.py +1 -0
- dbt_redis_adbc-0.1.0/dbt/adapters/redis_adbc/connections.py +286 -0
- dbt_redis_adbc-0.1.0/dbt/adapters/redis_adbc/impl.py +482 -0
- dbt_redis_adbc-0.1.0/dbt/include/redis_adbc/__init__.py +3 -0
- dbt_redis_adbc-0.1.0/dbt/include/redis_adbc/dbt_project.yml +5 -0
- dbt_redis_adbc-0.1.0/dbt/include/redis_adbc/macros/adapters.sql +72 -0
- dbt_redis_adbc-0.1.0/dbt/include/redis_adbc/macros/contracts.sql +26 -0
- dbt_redis_adbc-0.1.0/dbt/include/redis_adbc/macros/cross_db.sql +18 -0
- dbt_redis_adbc-0.1.0/dbt/include/redis_adbc/macros/docs_and_grants.sql +27 -0
- dbt_redis_adbc-0.1.0/dbt/include/redis_adbc/macros/materializations.sql +132 -0
- dbt_redis_adbc-0.1.0/dbt/include/redis_adbc/macros/python.sql +16 -0
- dbt_redis_adbc-0.1.0/pyproject.toml +33 -0
- dbt_redis_adbc-0.1.0/tests/functional/conftest.py +78 -0
- dbt_redis_adbc-0.1.0/tests/functional/test_basic.py +108 -0
- dbt_redis_adbc-0.1.0/tests/functional/test_caching.py +22 -0
- dbt_redis_adbc-0.1.0/tests/functional/test_clone.py +28 -0
- dbt_redis_adbc-0.1.0/tests/functional/test_concurrency.py +5 -0
- dbt_redis_adbc-0.1.0/tests/functional/test_constraints.py +217 -0
- dbt_redis_adbc-0.1.0/tests/functional/test_empty.py +13 -0
- dbt_redis_adbc-0.1.0/tests/functional/test_ephemeral.py +66 -0
- dbt_redis_adbc-0.1.0/tests/functional/test_hooks.py +132 -0
- dbt_redis_adbc-0.1.0/tests/functional/test_incremental.py +45 -0
- dbt_redis_adbc-0.1.0/tests/functional/test_not_supported.py +50 -0
- dbt_redis_adbc-0.1.0/tests/functional/test_persist_docs.py +32 -0
- dbt_redis_adbc-0.1.0/tests/functional/test_python_model.py +190 -0
- dbt_redis_adbc-0.1.0/tests/functional/test_query_comment.py +32 -0
- dbt_redis_adbc-0.1.0/tests/functional/test_retry.py +33 -0
- dbt_redis_adbc-0.1.0/tests/functional/test_sample_mode.py +5 -0
- dbt_redis_adbc-0.1.0/tests/functional/test_show.py +14 -0
- dbt_redis_adbc-0.1.0/tests/functional/test_snapshot.py +73 -0
- dbt_redis_adbc-0.1.0/tests/functional/test_store_test_failures.py +44 -0
- dbt_redis_adbc-0.1.0/tests/functional/test_unit_testing.py +29 -0
- dbt_redis_adbc-0.1.0/tests/functional/test_utils.py +286 -0
|
@@ -0,0 +1,88 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: dbt-redis-adbc
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: dbt adapter for Redis 8.x using the Redis ADBC driver
|
|
5
|
+
Project-URL: Homepage, https://github.com/alberttwong/redis-dbt-adapter-with-project
|
|
6
|
+
Project-URL: Documentation, https://github.com/alberttwong/redis-dbt-adapter-with-project#adapter-reference
|
|
7
|
+
Project-URL: Issues, https://github.com/alberttwong/redis-dbt-adapter-with-project/issues
|
|
8
|
+
Project-URL: Source, https://github.com/alberttwong/redis-dbt-adapter-with-project/tree/main/adapter
|
|
9
|
+
Author: Albert T. Wong
|
|
10
|
+
Keywords: adbc,analytics,dbt,redis
|
|
11
|
+
Classifier: Development Status :: 3 - Alpha
|
|
12
|
+
Classifier: Programming Language :: Python :: 3
|
|
13
|
+
Classifier: Programming Language :: Python :: 3 :: Only
|
|
14
|
+
Classifier: Topic :: Database
|
|
15
|
+
Requires-Python: >=3.10
|
|
16
|
+
Requires-Dist: adbc-driver-manager>=1.6
|
|
17
|
+
Requires-Dist: dbt-adapters<2,>=1.16
|
|
18
|
+
Requires-Dist: dbt-core<2,>=1.10
|
|
19
|
+
Requires-Dist: pyarrow>=18
|
|
20
|
+
Description-Content-Type: text/markdown
|
|
21
|
+
|
|
22
|
+
# dbt-redis-adbc
|
|
23
|
+
|
|
24
|
+
A community dbt adapter for Redis 8.x using the
|
|
25
|
+
[Redis Arrow Database Connectivity (ADBC) driver](https://github.com/alberttwong/redis-adbc-driver).
|
|
26
|
+
Supports Python-based dbt 1.x; dbt 2.x and the dbt platform are not supported.
|
|
27
|
+
|
|
28
|
+
## Install
|
|
29
|
+
|
|
30
|
+
Requires Python 3.10 or later and Redis 8.x with the Query Engine (Search)
|
|
31
|
+
enabled. The adapter is tested against Redis 8.6.2.
|
|
32
|
+
|
|
33
|
+
```shell
|
|
34
|
+
python -m pip install dbt-redis-adbc
|
|
35
|
+
```
|
|
36
|
+
|
|
37
|
+
The Python package does not include the native Redis ADBC driver. Download and
|
|
38
|
+
extract the shared library for your platform from the
|
|
39
|
+
[driver releases](https://github.com/alberttwong/redis-adbc-driver/releases).
|
|
40
|
+
The example project uses driver v0.0.14. Prebuilt libraries are available for
|
|
41
|
+
macOS Apple Silicon, Linux x86-64, and Linux arm64.
|
|
42
|
+
|
|
43
|
+
Set `REDIS_ADBC_DRIVER` to the library's absolute path. See the driver's
|
|
44
|
+
[server requirements and ACL permissions](https://github.com/alberttwong/redis-adbc-driver#server-requirements)
|
|
45
|
+
before connecting.
|
|
46
|
+
|
|
47
|
+
## Configure
|
|
48
|
+
|
|
49
|
+
Set your dbt project's `profile` to `redis_project` and add this profile to
|
|
50
|
+
`~/.dbt/profiles.yml`:
|
|
51
|
+
|
|
52
|
+
```yaml
|
|
53
|
+
redis_project:
|
|
54
|
+
target: dev
|
|
55
|
+
outputs:
|
|
56
|
+
dev:
|
|
57
|
+
type: redis_adbc
|
|
58
|
+
driver: "{{ env_var('REDIS_ADBC_DRIVER') }}"
|
|
59
|
+
uri: "{{ env_var('REDIS_URI', 'redis://localhost:6379/0') }}"
|
|
60
|
+
username: "{{ env_var('REDIS_USERNAME', '') }}"
|
|
61
|
+
password: "{{ env_var('REDIS_PASSWORD', '') }}"
|
|
62
|
+
database: redis
|
|
63
|
+
schema: analytics
|
|
64
|
+
threads: 4
|
|
65
|
+
```
|
|
66
|
+
|
|
67
|
+
Use a `rediss://` URI for TLS. Keep credentials in the separate username and
|
|
68
|
+
password fields. From your project, verify the connection with:
|
|
69
|
+
|
|
70
|
+
```shell
|
|
71
|
+
dbt debug --profile redis_project
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
## Features and limitations
|
|
75
|
+
|
|
76
|
+
- Tables, views, ephemeral models, Arrow bulk-loaded seeds, and snapshots
|
|
77
|
+
- SQL incremental models with `append`, `delete+insert`, `merge`, and `microbatch`
|
|
78
|
+
- Model contracts, persisted relation and column descriptions, and Python table models
|
|
79
|
+
- No SQL transactions or rollback: statements autocommit, and interrupted writes
|
|
80
|
+
may need cleanup or a rerun
|
|
81
|
+
- No materialized views or SQL grants; configure access with Redis ACLs
|
|
82
|
+
- Python models run in the dbt process and load their inputs into its memory;
|
|
83
|
+
incremental Python models are not supported
|
|
84
|
+
|
|
85
|
+
See the [adapter reference](https://github.com/alberttwong/redis-dbt-adapter-with-project#adapter-reference)
|
|
86
|
+
for all profile options, supported features, and limitations. The
|
|
87
|
+
[example project and functional tests](https://github.com/alberttwong/redis-dbt-adapter-with-project)
|
|
88
|
+
demonstrate the adapter with NYC taxi data.
|
|
@@ -0,0 +1,67 @@
|
|
|
1
|
+
# dbt-redis-adbc
|
|
2
|
+
|
|
3
|
+
A community dbt adapter for Redis 8.x using the
|
|
4
|
+
[Redis Arrow Database Connectivity (ADBC) driver](https://github.com/alberttwong/redis-adbc-driver).
|
|
5
|
+
Supports Python-based dbt 1.x; dbt 2.x and the dbt platform are not supported.
|
|
6
|
+
|
|
7
|
+
## Install
|
|
8
|
+
|
|
9
|
+
Requires Python 3.10 or later and Redis 8.x with the Query Engine (Search)
|
|
10
|
+
enabled. The adapter is tested against Redis 8.6.2.
|
|
11
|
+
|
|
12
|
+
```shell
|
|
13
|
+
python -m pip install dbt-redis-adbc
|
|
14
|
+
```
|
|
15
|
+
|
|
16
|
+
The Python package does not include the native Redis ADBC driver. Download and
|
|
17
|
+
extract the shared library for your platform from the
|
|
18
|
+
[driver releases](https://github.com/alberttwong/redis-adbc-driver/releases).
|
|
19
|
+
The example project uses driver v0.0.14. Prebuilt libraries are available for
|
|
20
|
+
macOS Apple Silicon, Linux x86-64, and Linux arm64.
|
|
21
|
+
|
|
22
|
+
Set `REDIS_ADBC_DRIVER` to the library's absolute path. See the driver's
|
|
23
|
+
[server requirements and ACL permissions](https://github.com/alberttwong/redis-adbc-driver#server-requirements)
|
|
24
|
+
before connecting.
|
|
25
|
+
|
|
26
|
+
## Configure
|
|
27
|
+
|
|
28
|
+
Set your dbt project's `profile` to `redis_project` and add this profile to
|
|
29
|
+
`~/.dbt/profiles.yml`:
|
|
30
|
+
|
|
31
|
+
```yaml
|
|
32
|
+
redis_project:
|
|
33
|
+
target: dev
|
|
34
|
+
outputs:
|
|
35
|
+
dev:
|
|
36
|
+
type: redis_adbc
|
|
37
|
+
driver: "{{ env_var('REDIS_ADBC_DRIVER') }}"
|
|
38
|
+
uri: "{{ env_var('REDIS_URI', 'redis://localhost:6379/0') }}"
|
|
39
|
+
username: "{{ env_var('REDIS_USERNAME', '') }}"
|
|
40
|
+
password: "{{ env_var('REDIS_PASSWORD', '') }}"
|
|
41
|
+
database: redis
|
|
42
|
+
schema: analytics
|
|
43
|
+
threads: 4
|
|
44
|
+
```
|
|
45
|
+
|
|
46
|
+
Use a `rediss://` URI for TLS. Keep credentials in the separate username and
|
|
47
|
+
password fields. From your project, verify the connection with:
|
|
48
|
+
|
|
49
|
+
```shell
|
|
50
|
+
dbt debug --profile redis_project
|
|
51
|
+
```
|
|
52
|
+
|
|
53
|
+
## Features and limitations
|
|
54
|
+
|
|
55
|
+
- Tables, views, ephemeral models, Arrow bulk-loaded seeds, and snapshots
|
|
56
|
+
- SQL incremental models with `append`, `delete+insert`, `merge`, and `microbatch`
|
|
57
|
+
- Model contracts, persisted relation and column descriptions, and Python table models
|
|
58
|
+
- No SQL transactions or rollback: statements autocommit, and interrupted writes
|
|
59
|
+
may need cleanup or a rerun
|
|
60
|
+
- No materialized views or SQL grants; configure access with Redis ACLs
|
|
61
|
+
- Python models run in the dbt process and load their inputs into its memory;
|
|
62
|
+
incremental Python models are not supported
|
|
63
|
+
|
|
64
|
+
See the [adapter reference](https://github.com/alberttwong/redis-dbt-adapter-with-project#adapter-reference)
|
|
65
|
+
for all profile options, supported features, and limitations. The
|
|
66
|
+
[example project and functional tests](https://github.com/alberttwong/redis-dbt-adapter-with-project)
|
|
67
|
+
demonstrate the adapter with NYC taxi data.
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
from dbt.adapters.base import AdapterPlugin
|
|
2
|
+
|
|
3
|
+
from dbt.adapters.redis_adbc.connections import RedisAdbcConnectionManager, RedisAdbcCredentials
|
|
4
|
+
from dbt.adapters.redis_adbc.impl import RedisAdbcAdapter
|
|
5
|
+
from dbt.include import redis_adbc
|
|
6
|
+
|
|
7
|
+
Plugin = AdapterPlugin(
|
|
8
|
+
adapter=RedisAdbcAdapter,
|
|
9
|
+
credentials=RedisAdbcCredentials,
|
|
10
|
+
include_path=redis_adbc.PACKAGE_PATH,
|
|
11
|
+
)
|
|
12
|
+
|
|
13
|
+
__all__ = ["Plugin", "RedisAdbcAdapter", "RedisAdbcConnectionManager", "RedisAdbcCredentials"]
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
version = "0.1.0"
|
|
@@ -0,0 +1,286 @@
|
|
|
1
|
+
"""Connection management: dbt <-> Redis ADBC driver via adbc_driver_manager's DB-API."""
|
|
2
|
+
|
|
3
|
+
import os
|
|
4
|
+
import sys
|
|
5
|
+
from contextlib import contextmanager
|
|
6
|
+
from dataclasses import dataclass
|
|
7
|
+
from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union
|
|
8
|
+
from urllib.parse import urlsplit, urlunsplit
|
|
9
|
+
|
|
10
|
+
import adbc_driver_manager
|
|
11
|
+
import adbc_driver_manager.dbapi as dbapi
|
|
12
|
+
from adbc_driver_manager import AdbcStatusCode
|
|
13
|
+
from dbt_common.exceptions import DbtDatabaseError, DbtRuntimeError
|
|
14
|
+
|
|
15
|
+
from dbt.adapters.contracts.connection import AdapterResponse, Connection, ConnectionState, Credentials
|
|
16
|
+
from dbt.adapters.events.logging import AdapterLogger
|
|
17
|
+
from dbt.adapters.sql import SQLConnectionManager
|
|
18
|
+
|
|
19
|
+
logger = AdapterLogger("RedisAdbc")
|
|
20
|
+
|
|
21
|
+
# Driver options that profile fields set (option -> field), so
|
|
22
|
+
# driver_options can't.
|
|
23
|
+
_PROFILE_OPTIONS = {
|
|
24
|
+
"uri": "uri",
|
|
25
|
+
"username": "username",
|
|
26
|
+
"password": "password",
|
|
27
|
+
"adbc.redis.default_schema": "schema",
|
|
28
|
+
"adbc.redis.aggregate_pushdown": "aggregate_pushdown",
|
|
29
|
+
"adbc.redis.rename_rekey": "rename_rekey",
|
|
30
|
+
"adbc.redis.time_zone": "time_zone",
|
|
31
|
+
"adbc.redis.read_timeout": "read_timeout",
|
|
32
|
+
"adbc.redis.write_timeout": "write_timeout",
|
|
33
|
+
}
|
|
34
|
+
|
|
35
|
+
# Unknown profile keys already warned about (profiles can load more than once).
|
|
36
|
+
_warned_keys: Set[str] = set()
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def _option_value(value: Any) -> str:
|
|
40
|
+
# ADBC options are strings; booleans as the driver spells them.
|
|
41
|
+
if isinstance(value, bool):
|
|
42
|
+
return "true" if value else "false"
|
|
43
|
+
return str(value)
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
@dataclass
|
|
47
|
+
class RedisAdbcCredentials(Credentials):
|
|
48
|
+
# Path to libadbc_driver_redis.{so,dylib,dll}; falls back to $REDIS_ADBC_DRIVER.
|
|
49
|
+
driver: Optional[str] = None
|
|
50
|
+
uri: str = "redis://localhost:6379/0"
|
|
51
|
+
username: Optional[str] = None
|
|
52
|
+
password: Optional[str] = None
|
|
53
|
+
# exact | all | none (see the driver README).
|
|
54
|
+
aggregate_pushdown: Optional[str] = None
|
|
55
|
+
# Move a renamed table's rows to its new name's keys (driver option
|
|
56
|
+
# adbc.redis.rename_rekey). dbt renames every table it builds, so this
|
|
57
|
+
# copies every row once more per build, in exchange for key names that
|
|
58
|
+
# match the table's.
|
|
59
|
+
rename_rekey: Optional[bool] = None
|
|
60
|
+
# The session time zone (adbc.redis.time_zone): what current_timestamp's
|
|
61
|
+
# text, localtimestamp and TIMESTAMP <-> TIMESTAMP WITH TIME ZONE
|
|
62
|
+
# conversions use. UTC by default.
|
|
63
|
+
time_zone: Optional[str] = None
|
|
64
|
+
# How long the client waits for each reply, and to send each command:
|
|
65
|
+
# "30s", "10m", a number of seconds, or 0 for no timeout.
|
|
66
|
+
read_timeout: Optional[Union[str, int]] = None
|
|
67
|
+
write_timeout: Optional[Union[str, int]] = None
|
|
68
|
+
# Any other driver option, passed as is (see the driver README's Options).
|
|
69
|
+
driver_options: Optional[Dict[str, Any]] = None
|
|
70
|
+
|
|
71
|
+
@classmethod
|
|
72
|
+
def __pre_deserialize__(cls, data):
|
|
73
|
+
data = super().__pre_deserialize__(data)
|
|
74
|
+
# The driver exposes a single catalog named "redis".
|
|
75
|
+
data.setdefault("database", "redis")
|
|
76
|
+
# dbt accepts profile keys the adapter doesn't have, and drops them.
|
|
77
|
+
for key in sorted(data.keys() - cls.__dataclass_fields__.keys() - _warned_keys):
|
|
78
|
+
_warned_keys.add(key)
|
|
79
|
+
logger.warning(f"`{key}` isn't a redis_adbc profile option; it's ignored")
|
|
80
|
+
for option, value in (data.get("driver_options") or {}).items():
|
|
81
|
+
if option in _PROFILE_OPTIONS:
|
|
82
|
+
raise DbtRuntimeError(
|
|
83
|
+
f"driver_options can't set {option}; use the profile's `{_PROFILE_OPTIONS[option]}`"
|
|
84
|
+
)
|
|
85
|
+
if not isinstance(value, (str, int, float, bool)):
|
|
86
|
+
raise DbtRuntimeError(f"driver_options: {option} must be a string, number or boolean")
|
|
87
|
+
return data
|
|
88
|
+
|
|
89
|
+
@classmethod
|
|
90
|
+
def validate(cls, data):
|
|
91
|
+
# dbt validates the profile before __pre_deserialize__ sets the default.
|
|
92
|
+
super().validate({"database": "redis", **data})
|
|
93
|
+
|
|
94
|
+
@property
|
|
95
|
+
def type(self) -> str:
|
|
96
|
+
return "redis_adbc"
|
|
97
|
+
|
|
98
|
+
@property
|
|
99
|
+
def unique_field(self) -> str:
|
|
100
|
+
# Hashed for dbt's anonymous usage stats: the server, not credentials.
|
|
101
|
+
return _redact_uri(self.uri, keep_user=False)
|
|
102
|
+
|
|
103
|
+
def _connection_keys(self) -> Tuple[str, ...]:
|
|
104
|
+
return (
|
|
105
|
+
"driver",
|
|
106
|
+
"uri",
|
|
107
|
+
"database",
|
|
108
|
+
"schema",
|
|
109
|
+
"username",
|
|
110
|
+
"aggregate_pushdown",
|
|
111
|
+
"rename_rekey",
|
|
112
|
+
"time_zone",
|
|
113
|
+
"read_timeout",
|
|
114
|
+
"write_timeout",
|
|
115
|
+
"driver_options",
|
|
116
|
+
)
|
|
117
|
+
|
|
118
|
+
def connection_info(self, *, with_aliases: bool = False) -> Iterable[Tuple[str, Any]]:
|
|
119
|
+
# dbt debug prints these and logs them: never the URI's password.
|
|
120
|
+
for key, value in super().connection_info(with_aliases=with_aliases):
|
|
121
|
+
yield key, _redact_uri(value) if key == "uri" and value else value
|
|
122
|
+
|
|
123
|
+
def driver_path(self) -> str:
|
|
124
|
+
path = self.driver or os.environ.get("REDIS_ADBC_DRIVER")
|
|
125
|
+
if not path:
|
|
126
|
+
raise DbtRuntimeError(
|
|
127
|
+
"No Redis ADBC driver configured: set `driver` in profiles.yml or $REDIS_ADBC_DRIVER"
|
|
128
|
+
)
|
|
129
|
+
path = os.path.expanduser(path)
|
|
130
|
+
# `driver/libadbc_driver_redis` means the platform's shared library.
|
|
131
|
+
if not os.path.splitext(os.path.basename(path))[1]:
|
|
132
|
+
path += {"darwin": ".dylib", "win32": ".dll"}.get(sys.platform, ".so")
|
|
133
|
+
return path
|
|
134
|
+
|
|
135
|
+
def db_kwargs(self) -> Dict[str, str]:
|
|
136
|
+
"""The driver's database options for this profile."""
|
|
137
|
+
kwargs = {option: _option_value(v) for option, v in (self.driver_options or {}).items()}
|
|
138
|
+
for option, name in _PROFILE_OPTIONS.items():
|
|
139
|
+
value = getattr(self, name)
|
|
140
|
+
# None, or "" from an unset env_var, leaves the driver's default.
|
|
141
|
+
if value is not None and value != "":
|
|
142
|
+
kwargs[option] = _option_value(value)
|
|
143
|
+
return kwargs
|
|
144
|
+
|
|
145
|
+
|
|
146
|
+
def _redact_uri(uri: str, keep_user: bool = True) -> str:
|
|
147
|
+
"""The URI with its password replaced by ****, or (keep_user=False)
|
|
148
|
+
without any credentials."""
|
|
149
|
+
try:
|
|
150
|
+
parts = urlsplit(uri)
|
|
151
|
+
host = parts.hostname or ""
|
|
152
|
+
if parts.port:
|
|
153
|
+
host = f"{host}:{parts.port}"
|
|
154
|
+
except ValueError:
|
|
155
|
+
return "<unparseable uri>"
|
|
156
|
+
if keep_user and (parts.username or parts.password):
|
|
157
|
+
user = parts.username or ""
|
|
158
|
+
host = f"{user}:****@{host}" if parts.password else f"{user}@{host}"
|
|
159
|
+
return urlunsplit((parts.scheme, host, parts.path, parts.query, parts.fragment))
|
|
160
|
+
|
|
161
|
+
|
|
162
|
+
class RedisAdbcHandle:
|
|
163
|
+
"""Wraps an ADBC DB-API connection so cursors dbt never closes get closed.
|
|
164
|
+
|
|
165
|
+
ADBC refuses to close a connection while it still has open statements.
|
|
166
|
+
"""
|
|
167
|
+
|
|
168
|
+
def __init__(self, conn: dbapi.Connection):
|
|
169
|
+
self.conn = conn
|
|
170
|
+
self._cursors: list = []
|
|
171
|
+
|
|
172
|
+
def cursor(self):
|
|
173
|
+
# Close finished cursors eagerly so statements don't pile up.
|
|
174
|
+
self._cursors = [c for c in self._cursors if not c._closed]
|
|
175
|
+
cur = self.conn.cursor()
|
|
176
|
+
self._cursors.append(cur)
|
|
177
|
+
return cur
|
|
178
|
+
|
|
179
|
+
def commit(self):
|
|
180
|
+
pass # autocommit only
|
|
181
|
+
|
|
182
|
+
def rollback(self):
|
|
183
|
+
pass # no transactions
|
|
184
|
+
|
|
185
|
+
def cancel(self):
|
|
186
|
+
"""Cancel the statement running on this connection, or the streamed
|
|
187
|
+
result it is reading (from another thread). Either then fails with
|
|
188
|
+
CANCELLED (see exception_handler).
|
|
189
|
+
|
|
190
|
+
A cursor with neither refuses with INVALID_STATE.
|
|
191
|
+
"""
|
|
192
|
+
for cur in list(self._cursors):
|
|
193
|
+
if cur._closed:
|
|
194
|
+
continue
|
|
195
|
+
try:
|
|
196
|
+
cur.adbc_cancel()
|
|
197
|
+
except adbc_driver_manager.Error as e:
|
|
198
|
+
if e.status_code != AdbcStatusCode.INVALID_STATE:
|
|
199
|
+
logger.warning(f"Redis ADBC: cancelling a statement failed: {e}")
|
|
200
|
+
|
|
201
|
+
def close(self):
|
|
202
|
+
for cur in self._cursors:
|
|
203
|
+
try:
|
|
204
|
+
cur.close()
|
|
205
|
+
except Exception:
|
|
206
|
+
pass
|
|
207
|
+
self._cursors = []
|
|
208
|
+
self.conn.close()
|
|
209
|
+
|
|
210
|
+
|
|
211
|
+
class RedisAdbcSession:
|
|
212
|
+
"""What a Python model gets as `session`: the thread's ADBC DB-API
|
|
213
|
+
connection, except that its cursors go through the handle (so dbt
|
|
214
|
+
closes them, and Ctrl-C cancels them) and it can't be closed."""
|
|
215
|
+
|
|
216
|
+
def __init__(self, handle: RedisAdbcHandle):
|
|
217
|
+
self._handle = handle
|
|
218
|
+
|
|
219
|
+
def cursor(self):
|
|
220
|
+
return self._handle.cursor()
|
|
221
|
+
|
|
222
|
+
def close(self):
|
|
223
|
+
raise DbtRuntimeError("A Python model's session is dbt's connection; it can't be closed")
|
|
224
|
+
|
|
225
|
+
def __getattr__(self, name):
|
|
226
|
+
return getattr(self._handle.conn, name)
|
|
227
|
+
|
|
228
|
+
|
|
229
|
+
class RedisAdbcConnectionManager(SQLConnectionManager):
|
|
230
|
+
TYPE = "redis_adbc"
|
|
231
|
+
|
|
232
|
+
@classmethod
|
|
233
|
+
def open(cls, connection: Connection) -> Connection:
|
|
234
|
+
if connection.state == ConnectionState.OPEN:
|
|
235
|
+
return connection
|
|
236
|
+
creds: RedisAdbcCredentials = connection.credentials
|
|
237
|
+
|
|
238
|
+
def connect():
|
|
239
|
+
conn = dbapi.connect(driver=creds.driver_path(), db_kwargs=creds.db_kwargs(), autocommit=True)
|
|
240
|
+
return RedisAdbcHandle(conn)
|
|
241
|
+
|
|
242
|
+
return cls.retry_connection(
|
|
243
|
+
connection,
|
|
244
|
+
connect=connect,
|
|
245
|
+
logger=logger,
|
|
246
|
+
retryable_exceptions=[adbc_driver_manager.OperationalError],
|
|
247
|
+
)
|
|
248
|
+
|
|
249
|
+
@contextmanager
|
|
250
|
+
def exception_handler(self, sql: str):
|
|
251
|
+
try:
|
|
252
|
+
yield
|
|
253
|
+
except adbc_driver_manager.Error as e:
|
|
254
|
+
if e.status_code == AdbcStatusCode.CANCELLED:
|
|
255
|
+
# Cancelled (Ctrl-C, --fail-fast): not the database's error.
|
|
256
|
+
# A write keeps the rows it wrote before the cancel.
|
|
257
|
+
logger.debug(f"Redis ADBC: cancelled: {e}")
|
|
258
|
+
raise DbtRuntimeError(str(e).strip()) from e
|
|
259
|
+
logger.debug(f"Redis ADBC error: {e}")
|
|
260
|
+
raise DbtDatabaseError(str(e).strip()) from e
|
|
261
|
+
except Exception as e:
|
|
262
|
+
if isinstance(e, (DbtRuntimeError, DbtDatabaseError)):
|
|
263
|
+
raise
|
|
264
|
+
raise DbtRuntimeError(str(e)) from e
|
|
265
|
+
|
|
266
|
+
def execute(self, sql: str, auto_begin: bool = False, fetch: bool = False, limit: Optional[int] = None):
|
|
267
|
+
# The driver streams a long result (over 65,536 rows), so an error
|
|
268
|
+
# past its first batch comes from fetching the rows, which dbt does
|
|
269
|
+
# outside add_query's exception_handler.
|
|
270
|
+
with self.exception_handler(sql):
|
|
271
|
+
return super().execute(sql, auto_begin=auto_begin, fetch=fetch, limit=limit)
|
|
272
|
+
|
|
273
|
+
def cancel(self, connection: Connection):
|
|
274
|
+
# dbt calls this from the main thread on Ctrl-C and --fail-fast.
|
|
275
|
+
connection.handle.cancel()
|
|
276
|
+
|
|
277
|
+
@classmethod
|
|
278
|
+
def get_response(cls, cursor: Any) -> AdapterResponse:
|
|
279
|
+
rows = cursor.rowcount if cursor is not None else -1
|
|
280
|
+
return AdapterResponse(_message="OK", rows_affected=rows if rows is not None else -1)
|
|
281
|
+
|
|
282
|
+
@classmethod
|
|
283
|
+
def data_type_code_to_name(cls, type_code) -> str:
|
|
284
|
+
from dbt.adapters.redis_adbc.impl import arrow_type_to_sql
|
|
285
|
+
|
|
286
|
+
return arrow_type_to_sql(type_code)
|