arrows 0.2.0__tar.gz → 0.2.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.
- {arrows-0.2.0 → arrows-0.2.2}/PKG-INFO +7 -2
- {arrows-0.2.0 → arrows-0.2.2}/README.md +5 -1
- {arrows-0.2.0 → arrows-0.2.2}/pyproject.toml +2 -1
- {arrows-0.2.0 → arrows-0.2.2}/pyproject.toml.orig +2 -1
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/components/aws.py +69 -6
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/components/s3.py +28 -17
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/s3.py +52 -12
- arrows-0.2.2/tests/test_s3_credentials.py +229 -0
- {arrows-0.2.0 → arrows-0.2.2}/LICENSE +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/examples/README.md +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/examples/bootstrap_aws.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/__init__.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/__main__.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/auth.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/cli.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/components/__init__.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/components/gmail.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/components/google.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/components/google_sheets.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/components/redshift.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/components/sqlite.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/core/__init__.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/core/component.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/core/config.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/core/engine.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/core/errors.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/core/registry.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/core/secrets.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/core/session.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/gmail.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/google_sheets.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/mysql.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/py.typed +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/redis.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/redshift.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/spark.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/sqlite.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/template_renderer.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/src/arrows/utils.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/tests/conftest.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/tests/test_login.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/tests/test_packaging.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/tests/test_public_api.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/tests/test_registry_and_session.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/tests/test_secrets.py +0 -0
- {arrows-0.2.0 → arrows-0.2.2}/tests/test_sqlite_component.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: arrows
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.2
|
|
4
4
|
Summary: Universal data wrangling tool built around Apache Arrow, DuckDB, Polars, S3 for efficiently transferring and transforming data between different data sources with pluggable, independently loadable components
|
|
5
5
|
Keywords: arrow,duckdb,polars,etl,s3,redshift,data-engineering
|
|
6
6
|
Author: George Jiang
|
|
@@ -28,6 +28,7 @@ Requires-Dist: google-auth>=2.58.0
|
|
|
28
28
|
Requires-Dist: pandas>=2.2
|
|
29
29
|
Requires-Dist: polars>=1.44.2
|
|
30
30
|
Requires-Dist: keyring>=25.0
|
|
31
|
+
Requires-Dist: redshift-connector>=2.1.16
|
|
31
32
|
Requires-Python: >=3.11
|
|
32
33
|
Project-URL: Homepage, https://github.com/George-Jiang/arrows
|
|
33
34
|
Project-URL: Repository, https://github.com/George-Jiang/arrows
|
|
@@ -614,10 +615,14 @@ sql = render_template('path/to/template.sql', table_name='my_table', date='2024-
|
|
|
614
615
|
```bash
|
|
615
616
|
uv sync # install with every component
|
|
616
617
|
uv run pytest # unit tests, no credentials needed
|
|
617
|
-
uv run ruff check
|
|
618
|
+
uv run ruff check . # lint
|
|
619
|
+
uv run ruff format . # format; CI checks this
|
|
618
620
|
uv run arrows doctor # check real connectivity for what is configured
|
|
619
621
|
```
|
|
620
622
|
|
|
623
|
+
Publishing a new version is
|
|
624
|
+
[docs/RELEASING.md](https://github.com/George-Jiang/arrows/blob/main/docs/RELEASING.md).
|
|
625
|
+
|
|
621
626
|
## Notes
|
|
622
627
|
|
|
623
628
|
1. **Load what you need**: `arrows.load(...)` fails fast on a missing credential; `arrows doctor` checks a real round-trip
|
|
@@ -577,10 +577,14 @@ sql = render_template('path/to/template.sql', table_name='my_table', date='2024-
|
|
|
577
577
|
```bash
|
|
578
578
|
uv sync # install with every component
|
|
579
579
|
uv run pytest # unit tests, no credentials needed
|
|
580
|
-
uv run ruff check
|
|
580
|
+
uv run ruff check . # lint
|
|
581
|
+
uv run ruff format . # format; CI checks this
|
|
581
582
|
uv run arrows doctor # check real connectivity for what is configured
|
|
582
583
|
```
|
|
583
584
|
|
|
585
|
+
Publishing a new version is
|
|
586
|
+
[docs/RELEASING.md](https://github.com/George-Jiang/arrows/blob/main/docs/RELEASING.md).
|
|
587
|
+
|
|
584
588
|
## Notes
|
|
585
589
|
|
|
586
590
|
1. **Load what you need**: `arrows.load(...)` fails fast on a missing credential; `arrows doctor` checks a real round-trip
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
[project]
|
|
2
2
|
name = "arrows"
|
|
3
|
-
version = "0.2.
|
|
3
|
+
version = "0.2.2"
|
|
4
4
|
description = "Universal data wrangling tool built around Apache Arrow, DuckDB, Polars, S3 for efficiently transferring and transforming data between different data sources with pluggable, independently loadable components"
|
|
5
5
|
readme = "README.md"
|
|
6
6
|
requires-python = ">=3.11"
|
|
@@ -39,6 +39,7 @@ dependencies = [
|
|
|
39
39
|
"pandas>=2.2",
|
|
40
40
|
"polars>=1.44.2",
|
|
41
41
|
"keyring>=25.0",
|
|
42
|
+
"redshift-connector>=2.1.16",
|
|
42
43
|
]
|
|
43
44
|
|
|
44
45
|
[[project.authors]]
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
[project]
|
|
2
2
|
name = "arrows"
|
|
3
|
-
version = "0.2.
|
|
3
|
+
version = "0.2.2"
|
|
4
4
|
description = "Universal data wrangling tool built around Apache Arrow, DuckDB, Polars, S3 for efficiently transferring and transforming data between different data sources with pluggable, independently loadable components"
|
|
5
5
|
readme = "README.md"
|
|
6
6
|
authors = [
|
|
@@ -47,6 +47,7 @@ dependencies = [
|
|
|
47
47
|
"polars>=1.44.2",
|
|
48
48
|
# secret storage for login(save=True)
|
|
49
49
|
"keyring>=25.0",
|
|
50
|
+
"redshift-connector>=2.1.16",
|
|
50
51
|
]
|
|
51
52
|
|
|
52
53
|
[project.urls]
|
|
@@ -73,6 +73,7 @@ class AwsComponent(Component):
|
|
|
73
73
|
self.boto3_session: Any = None
|
|
74
74
|
self.region: str | None = None
|
|
75
75
|
self._explicit_credentials = False
|
|
76
|
+
self._duckdb_credentials: tuple[str, str, str | None] | None = None
|
|
76
77
|
|
|
77
78
|
def setup(self, secrets: SecretStore) -> None:
|
|
78
79
|
boto3 = self.import_module('boto3')
|
|
@@ -104,6 +105,60 @@ class AwsComponent(Component):
|
|
|
104
105
|
self.boto3_session = boto3.Session(**kwargs)
|
|
105
106
|
self._register_duckdb_secret()
|
|
106
107
|
|
|
108
|
+
# -- credentials -------------------------------------------------------
|
|
109
|
+
def frozen_credentials(self) -> Any:
|
|
110
|
+
"""The credentials to use *right now*, or ``None`` if there are none.
|
|
111
|
+
|
|
112
|
+
Every consumer has to ask for these at the moment of use rather than
|
|
113
|
+
keeping a copy. botocore hands back a refreshable credential object for
|
|
114
|
+
SSO, assumed roles and instance profiles, and ``get_frozen_credentials``
|
|
115
|
+
is what triggers the renewal when the current ones are near expiry. A
|
|
116
|
+
caller that froze them once is holding keys that stop working an hour
|
|
117
|
+
later, with no way to notice.
|
|
118
|
+
"""
|
|
119
|
+
credentials = self.boto3_session.get_credentials() if self.boto3_session else None
|
|
120
|
+
return credentials.get_frozen_credentials() if credentials else None
|
|
121
|
+
|
|
122
|
+
def credential_provider(self):
|
|
123
|
+
"""A Polars credential provider bound to this component's identity.
|
|
124
|
+
|
|
125
|
+
Without one, Polars builds a *fresh* ``boto3.Session()`` of its own and
|
|
126
|
+
resolves credentials from scratch. That session cannot see the arrows
|
|
127
|
+
SecretStore, so a profile or key supplied through ``login()``, ``.env``
|
|
128
|
+
or the keychain is invisible to it and it silently falls back to the
|
|
129
|
+
default profile — which is how a working DuckDB read sits next to a
|
|
130
|
+
Polars read failing on an expired SSO token.
|
|
131
|
+
"""
|
|
132
|
+
|
|
133
|
+
def provider():
|
|
134
|
+
credentials = self.boto3_session.get_credentials() if self.boto3_session else None
|
|
135
|
+
if credentials is None:
|
|
136
|
+
raise RuntimeError('No AWS credentials resolved; cannot read from S3.')
|
|
137
|
+
# Freeze first: on a deferred credential object the expiry does not
|
|
138
|
+
# exist until the keys have actually been resolved.
|
|
139
|
+
frozen = credentials.get_frozen_credentials()
|
|
140
|
+
values = {'aws_access_key_id': frozen.access_key, 'aws_secret_access_key': frozen.secret_key}
|
|
141
|
+
if frozen.token is not None:
|
|
142
|
+
values['aws_session_token'] = frozen.token
|
|
143
|
+
expiry = getattr(credentials, '_expiry_time', None)
|
|
144
|
+
return values, int(expiry.timestamp()) if expiry is not None else None
|
|
145
|
+
|
|
146
|
+
return provider
|
|
147
|
+
|
|
148
|
+
def storage_options(self) -> dict[str, str]:
|
|
149
|
+
"""Non-credential object-store settings for Polars."""
|
|
150
|
+
return {'aws_region': self.region} if self.region else {}
|
|
151
|
+
|
|
152
|
+
# -- duckdb ------------------------------------------------------------
|
|
153
|
+
def prepare_duckdb(self) -> None:
|
|
154
|
+
"""Make sure DuckDB holds usable credentials before a query runs.
|
|
155
|
+
|
|
156
|
+
Cheap and idempotent while nothing has rotated, so callers can put it in
|
|
157
|
+
front of every S3 read and write.
|
|
158
|
+
"""
|
|
159
|
+
if self._explicit_credentials:
|
|
160
|
+
self._register_duckdb_secret()
|
|
161
|
+
|
|
107
162
|
def _register_duckdb_secret(self) -> None:
|
|
108
163
|
"""Give DuckDB the same identity boto3 resolved.
|
|
109
164
|
|
|
@@ -111,36 +166,43 @@ class AwsComponent(Component):
|
|
|
111
166
|
|
|
112
167
|
* When boto3 resolved credentials from its own chain (profile, SSO,
|
|
113
168
|
instance role), DuckDB is pointed at that same chain — no key material
|
|
114
|
-
is written into SQL, and DuckDB refreshes on its own.
|
|
169
|
+
is written into SQL, and DuckDB refreshes on its own. Registered once.
|
|
115
170
|
* When the credentials came from the arrows SecretStore, DuckDB cannot
|
|
116
171
|
see that store, so the keys have to be handed over as a literal
|
|
117
|
-
secret.
|
|
118
|
-
|
|
172
|
+
secret. Those are a snapshot, so the secret is re-registered whenever
|
|
173
|
+
the underlying credentials change. They are visible in
|
|
174
|
+
``duckdb_secrets()``, which is the cost of that setup and a reason to
|
|
175
|
+
prefer a real AWS profile.
|
|
119
176
|
|
|
120
177
|
A failure here is not fatal: boto3 and pyarrow still work.
|
|
121
178
|
"""
|
|
122
179
|
import warnings
|
|
123
180
|
|
|
124
|
-
|
|
125
|
-
if
|
|
181
|
+
frozen = self.frozen_credentials()
|
|
182
|
+
if frozen is None:
|
|
126
183
|
warnings.warn('No AWS credentials resolved; DuckDB S3 access is unconfigured.', stacklevel=2)
|
|
127
184
|
return
|
|
128
185
|
|
|
129
186
|
region = f', REGION {quote_literal(self.region)}' if self.region else ''
|
|
130
187
|
if self._explicit_credentials:
|
|
131
|
-
|
|
188
|
+
current = (frozen.access_key, frozen.secret_key, frozen.token)
|
|
189
|
+
if current == self._duckdb_credentials:
|
|
190
|
+
return
|
|
132
191
|
token = f', SESSION_TOKEN {quote_literal(frozen.token)}' if frozen.token else ''
|
|
133
192
|
body = (
|
|
134
193
|
f'TYPE s3, KEY_ID {quote_literal(frozen.access_key)}, '
|
|
135
194
|
f'SECRET {quote_literal(frozen.secret_key)}{token}{region}'
|
|
136
195
|
)
|
|
137
196
|
else:
|
|
197
|
+
current = None
|
|
138
198
|
body = f'TYPE s3, PROVIDER credential_chain{region}'
|
|
139
199
|
|
|
140
200
|
try:
|
|
141
201
|
get_duckdb().execute(f'CREATE OR REPLACE SECRET arrows_s3 ({body});')
|
|
142
202
|
except Exception as exc:
|
|
143
203
|
warnings.warn(f'Could not register the DuckDB S3 secret: {exc}', stacklevel=2)
|
|
204
|
+
return
|
|
205
|
+
self._duckdb_credentials = current
|
|
144
206
|
|
|
145
207
|
def client(self, service: str, **kwargs) -> Any:
|
|
146
208
|
return self.boto3_session.client(service, **kwargs)
|
|
@@ -159,6 +221,7 @@ class AwsComponent(Component):
|
|
|
159
221
|
with contextlib.suppress(Exception):
|
|
160
222
|
get_duckdb().execute('DROP SECRET IF EXISTS arrows_s3;')
|
|
161
223
|
self.boto3_session = None
|
|
224
|
+
self._duckdb_credentials = None
|
|
162
225
|
|
|
163
226
|
|
|
164
227
|
COMPONENT = AwsComponent
|
|
@@ -17,6 +17,7 @@ class S3Component(Component):
|
|
|
17
17
|
def __init__(self) -> None:
|
|
18
18
|
super().__init__()
|
|
19
19
|
self._filesystem: Any = None
|
|
20
|
+
self._filesystem_credentials: tuple[str, str, str | None] | None = None
|
|
20
21
|
self.default_bucket: str | None = None
|
|
21
22
|
|
|
22
23
|
def setup(self, secrets: SecretStore) -> None:
|
|
@@ -26,23 +27,32 @@ class S3Component(Component):
|
|
|
26
27
|
|
|
27
28
|
@property
|
|
28
29
|
def filesystem(self) -> Any:
|
|
29
|
-
"""
|
|
30
|
-
|
|
31
|
-
|
|
32
|
-
|
|
33
|
-
|
|
34
|
-
|
|
35
|
-
|
|
36
|
-
|
|
37
|
-
|
|
38
|
-
|
|
39
|
-
|
|
40
|
-
|
|
41
|
-
|
|
42
|
-
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
|
|
30
|
+
"""A :class:`pyarrow.fs.S3FileSystem` holding the current credentials.
|
|
31
|
+
|
|
32
|
+
pyarrow copies the keys at construction and never refreshes them, so a
|
|
33
|
+
filesystem cached for the life of the session stops working an hour into
|
|
34
|
+
an SSO or instance-role session. Rebuilding it when the credentials
|
|
35
|
+
change is the fix; constructing one is cheap and opens no connection.
|
|
36
|
+
"""
|
|
37
|
+
from pyarrow.fs import S3FileSystem
|
|
38
|
+
|
|
39
|
+
aws = self.session.get('aws')
|
|
40
|
+
frozen = aws.frozen_credentials()
|
|
41
|
+
current = (frozen.access_key, frozen.secret_key, frozen.token) if frozen else None
|
|
42
|
+
if self._filesystem is not None and current == self._filesystem_credentials:
|
|
43
|
+
return self._filesystem
|
|
44
|
+
|
|
45
|
+
kwargs: dict[str, Any] = {}
|
|
46
|
+
if frozen is not None:
|
|
47
|
+
kwargs = {
|
|
48
|
+
'access_key': frozen.access_key,
|
|
49
|
+
'secret_key': frozen.secret_key,
|
|
50
|
+
'session_token': frozen.token,
|
|
51
|
+
}
|
|
52
|
+
if aws.region:
|
|
53
|
+
kwargs['region'] = aws.region
|
|
54
|
+
self._filesystem = S3FileSystem(**kwargs)
|
|
55
|
+
self._filesystem_credentials = current
|
|
46
56
|
return self._filesystem
|
|
47
57
|
|
|
48
58
|
def health_check(self) -> HealthStatus:
|
|
@@ -56,6 +66,7 @@ class S3Component(Component):
|
|
|
56
66
|
|
|
57
67
|
def close(self) -> None:
|
|
58
68
|
self._filesystem = None
|
|
69
|
+
self._filesystem_credentials = None
|
|
59
70
|
|
|
60
71
|
|
|
61
72
|
COMPONENT = S3Component
|
|
@@ -30,6 +30,20 @@ def _component() -> Any:
|
|
|
30
30
|
return _get_component('s3')
|
|
31
31
|
|
|
32
32
|
|
|
33
|
+
def _aws() -> Any:
|
|
34
|
+
return _get_component('aws')
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def _duckdb() -> Any:
|
|
38
|
+
"""The DuckDB connection, with its S3 credentials known to be current.
|
|
39
|
+
|
|
40
|
+
Mirrors ``arrows.google_sheets._duckdb``: credentials are refreshed at the
|
|
41
|
+
moment of use, never handed over once at load time.
|
|
42
|
+
"""
|
|
43
|
+
_aws().prepare_duckdb()
|
|
44
|
+
return get_duckdb()
|
|
45
|
+
|
|
46
|
+
|
|
33
47
|
def set_default_bucket_name(default_bucket_name: str) -> None:
|
|
34
48
|
"""Set the bucket used when a dataset is created without an explicit path."""
|
|
35
49
|
from .core.session import default_session
|
|
@@ -89,6 +103,27 @@ class S3Dataset:
|
|
|
89
103
|
"""The pyarrow S3 filesystem owned by the ``s3`` component."""
|
|
90
104
|
return _component().filesystem
|
|
91
105
|
|
|
106
|
+
@property
|
|
107
|
+
def _key(self) -> str:
|
|
108
|
+
"""The path without its scheme.
|
|
109
|
+
|
|
110
|
+
pyarrow does not strip ``s3://`` when a filesystem is passed explicitly,
|
|
111
|
+
so handing it the URI would address a literal key named ``s3:``.
|
|
112
|
+
"""
|
|
113
|
+
return self.s3_path[5:]
|
|
114
|
+
|
|
115
|
+
def _polars_options(self) -> dict[str, Any]:
|
|
116
|
+
"""Credentials for Polars, taken from the ``aws`` component.
|
|
117
|
+
|
|
118
|
+
Polars otherwise builds its own ``boto3.Session()``, which cannot see
|
|
119
|
+
anything supplied through ``login()``, ``.env`` or the keychain.
|
|
120
|
+
"""
|
|
121
|
+
aws = _aws()
|
|
122
|
+
return {
|
|
123
|
+
'credential_provider': aws.credential_provider(),
|
|
124
|
+
'storage_options': aws.storage_options() or None,
|
|
125
|
+
}
|
|
126
|
+
|
|
92
127
|
def __repr__(self) -> str:
|
|
93
128
|
return f'S3Dataset: {self.s3_path}'
|
|
94
129
|
|
|
@@ -97,16 +132,17 @@ class S3Dataset:
|
|
|
97
132
|
if engine == 'pyarrow':
|
|
98
133
|
import pyarrow.parquet as pq
|
|
99
134
|
|
|
100
|
-
return pq.read_table(self.
|
|
135
|
+
return pq.read_table(self._key, filesystem=self.s3)
|
|
101
136
|
return self.to_duckdb().to_arrow_table()
|
|
102
137
|
|
|
103
138
|
def to_duckdb(self):
|
|
104
|
-
return
|
|
139
|
+
return _duckdb().from_parquet(f'{self.s3_path}*.parquet')
|
|
105
140
|
|
|
106
141
|
def to_polars(self, lazy: bool = False):
|
|
107
142
|
import polars as pl
|
|
108
143
|
|
|
109
|
-
|
|
144
|
+
options = self._polars_options()
|
|
145
|
+
return pl.scan_parquet(self.s3_path, **options) if lazy else pl.read_parquet(self.s3_path, **options)
|
|
110
146
|
|
|
111
147
|
# -- write -------------------------------------------------------------
|
|
112
148
|
def from_arrow(self, arrow, engine: str = 'pyarrow') -> None:
|
|
@@ -114,19 +150,23 @@ class S3Dataset:
|
|
|
114
150
|
if engine == 'pyarrow':
|
|
115
151
|
import pyarrow.parquet as pq
|
|
116
152
|
|
|
117
|
-
pq.write_to_dataset(arrow, self.
|
|
153
|
+
pq.write_to_dataset(arrow, self._key, filesystem=self.s3)
|
|
118
154
|
else:
|
|
119
|
-
|
|
155
|
+
_duckdb().execute(f"COPY arrow TO '{self.s3_path[:-1]}' (FORMAT parquet, FILE_SIZE_BYTES '1G')")
|
|
120
156
|
|
|
121
157
|
def from_polars(self, df) -> None:
|
|
122
158
|
import polars as pl
|
|
123
159
|
|
|
124
160
|
self.clear_contents()
|
|
125
|
-
|
|
161
|
+
# `PartitionMaxSize` was removed in the Polars version this package
|
|
162
|
+
# requires; `PartitionBy` without a key is the same thing — one file per
|
|
163
|
+
# 512k rows — and is what remains.
|
|
164
|
+
partition_info = pl.PartitionBy(self.s3_path, max_rows_per_file=512_000)
|
|
165
|
+
options = self._polars_options()
|
|
126
166
|
if isinstance(df, pl.LazyFrame):
|
|
127
|
-
df.sink_parquet(partition_info)
|
|
167
|
+
df.sink_parquet(partition_info, **options)
|
|
128
168
|
else:
|
|
129
|
-
df.write_parquet(partition_info)
|
|
169
|
+
df.write_parquet(partition_info, **options)
|
|
130
170
|
|
|
131
171
|
def from_redshift(self, sql: str, **kwargs) -> None:
|
|
132
172
|
from . import redshift
|
|
@@ -148,7 +188,7 @@ class S3Dataset:
|
|
|
148
188
|
source = quote_literal(f'{self.s3_path}*.parquet')
|
|
149
189
|
body = _parse_self_sql(rendered, 'self', 'temp_s3_dataset_table')
|
|
150
190
|
# `source` is a quoted literal; `body` is the caller's own SQL.
|
|
151
|
-
return
|
|
191
|
+
return _duckdb().sql(f'WITH temp_s3_dataset_table AS (SELECT * FROM {source})\n{body}') # noqa: S608
|
|
152
192
|
|
|
153
193
|
def sql(self, sql: str, **kwargs):
|
|
154
194
|
"""Run SQL against a lazily scanned Polars frame of this dataset."""
|
|
@@ -156,13 +196,13 @@ class S3Dataset:
|
|
|
156
196
|
|
|
157
197
|
df = self.to_polars(lazy=True) # noqa: F841 - referenced by DuckDB replacement scan
|
|
158
198
|
rendered = _parse_self_sql(Template(sql).render(**kwargs), 'self', 'df')
|
|
159
|
-
return
|
|
199
|
+
return _duckdb().sql(rendered)
|
|
160
200
|
|
|
161
201
|
# -- lifecycle ---------------------------------------------------------
|
|
162
202
|
def delete(self) -> None:
|
|
163
203
|
from pyarrow.fs import FileType
|
|
164
204
|
|
|
165
|
-
path = self.
|
|
205
|
+
path = self._key
|
|
166
206
|
file_type = self.s3.get_file_info(path).type
|
|
167
207
|
if file_type == FileType.Directory:
|
|
168
208
|
self.s3.delete_dir(path)
|
|
@@ -174,7 +214,7 @@ class S3Dataset:
|
|
|
174
214
|
def clear_contents(self) -> None:
|
|
175
215
|
from pyarrow.fs import FileType
|
|
176
216
|
|
|
177
|
-
path = self.
|
|
217
|
+
path = self._key
|
|
178
218
|
if self.s3.get_file_info(path).type == FileType.Directory:
|
|
179
219
|
self.s3.delete_dir_contents(path)
|
|
180
220
|
|
|
@@ -0,0 +1,229 @@
|
|
|
1
|
+
"""Every S3 path must use the credentials the ``aws`` component resolved.
|
|
2
|
+
|
|
3
|
+
The failure these guard against is quiet and asymmetric: a library left to find
|
|
4
|
+
its own credentials (Polars builds its own ``boto3.Session``, pyarrow its own
|
|
5
|
+
filesystem) cannot see the arrows SecretStore, so a DuckDB read succeeds while a
|
|
6
|
+
Polars read fails on a stale profile. The second half of the problem is time —
|
|
7
|
+
credentials handed over once stop working when they expire, so they have to be
|
|
8
|
+
fetched at the moment of use.
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
from __future__ import annotations
|
|
12
|
+
|
|
13
|
+
import pyarrow as pa
|
|
14
|
+
import pytest
|
|
15
|
+
|
|
16
|
+
from arrows import s3
|
|
17
|
+
from arrows.core.engine import reset_duckdb, set_duckdb
|
|
18
|
+
from arrows.core.secrets import MappingProvider, SecretStore
|
|
19
|
+
from arrows.core.session import Session
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class FakeFrozenCredentials:
|
|
23
|
+
def __init__(self, access_key: str, secret_key: str, token: str | None):
|
|
24
|
+
self.access_key = access_key
|
|
25
|
+
self.secret_key = secret_key
|
|
26
|
+
self.token = token
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
class FakeCredentials:
|
|
30
|
+
"""Stands in for botocore's refreshable credentials.
|
|
31
|
+
|
|
32
|
+
Like the real thing, the current keys are only visible through
|
|
33
|
+
``get_frozen_credentials()``, and they can change between calls.
|
|
34
|
+
"""
|
|
35
|
+
|
|
36
|
+
def __init__(self) -> None:
|
|
37
|
+
self.frozen = FakeFrozenCredentials('AKIA-first', 'secret-first', 'token-first')
|
|
38
|
+
self._expiry_time = None
|
|
39
|
+
self.calls = 0
|
|
40
|
+
|
|
41
|
+
def get_frozen_credentials(self) -> FakeFrozenCredentials:
|
|
42
|
+
self.calls += 1
|
|
43
|
+
return self.frozen
|
|
44
|
+
|
|
45
|
+
def rotate(self) -> None:
|
|
46
|
+
self.frozen = FakeFrozenCredentials('AKIA-second', 'secret-second', 'token-second')
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
class FakeBoto3Session:
|
|
50
|
+
def __init__(self, credentials: FakeCredentials):
|
|
51
|
+
self._credentials = credentials
|
|
52
|
+
|
|
53
|
+
def get_credentials(self) -> FakeCredentials:
|
|
54
|
+
return self._credentials
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
class FakeConnection:
|
|
58
|
+
"""A DuckDB stand-in that only records what was executed."""
|
|
59
|
+
|
|
60
|
+
def __init__(self) -> None:
|
|
61
|
+
self.statements: list[str] = []
|
|
62
|
+
|
|
63
|
+
def execute(self, sql: str):
|
|
64
|
+
self.statements.append(sql)
|
|
65
|
+
return self
|
|
66
|
+
|
|
67
|
+
def from_parquet(self, glob: str):
|
|
68
|
+
self.statements.append(f'read_parquet({glob})')
|
|
69
|
+
return self
|
|
70
|
+
|
|
71
|
+
def secrets(self) -> list[str]:
|
|
72
|
+
return [s for s in self.statements if 'SECRET arrows_s3' in s]
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
@pytest.fixture
|
|
76
|
+
def loaded(monkeypatch):
|
|
77
|
+
"""An s3 component whose AWS identity is a controllable fake."""
|
|
78
|
+
connection = FakeConnection()
|
|
79
|
+
set_duckdb(connection)
|
|
80
|
+
secrets = SecretStore(
|
|
81
|
+
[
|
|
82
|
+
MappingProvider(
|
|
83
|
+
{
|
|
84
|
+
'AWS_ACCESS_KEY_ID': 'AKIA-first',
|
|
85
|
+
'AWS_SECRET_ACCESS_KEY': 'secret-first',
|
|
86
|
+
'AWS_SESSION_TOKEN': 'token-first',
|
|
87
|
+
'AWS_REGION': 'us-east-1',
|
|
88
|
+
'ARROWS_DEFAULT_BUCKET': 'demo-bucket',
|
|
89
|
+
},
|
|
90
|
+
name='m',
|
|
91
|
+
)
|
|
92
|
+
]
|
|
93
|
+
)
|
|
94
|
+
session = Session(secrets=secrets)
|
|
95
|
+
credentials = FakeCredentials()
|
|
96
|
+
with session.activate():
|
|
97
|
+
session.load('s3')
|
|
98
|
+
session.get('aws').boto3_session = FakeBoto3Session(credentials)
|
|
99
|
+
yield session, credentials, connection
|
|
100
|
+
session.close()
|
|
101
|
+
reset_duckdb()
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
# -- Polars ----------------------------------------------------------------
|
|
105
|
+
def test_to_polars_passes_the_component_credentials(loaded, monkeypatch):
|
|
106
|
+
"""Without this, Polars resolves credentials from a session of its own."""
|
|
107
|
+
import polars as pl
|
|
108
|
+
|
|
109
|
+
captured: dict = {}
|
|
110
|
+
|
|
111
|
+
def fake_read_parquet(source, **kwargs):
|
|
112
|
+
captured.update(source=source, **kwargs)
|
|
113
|
+
return 'frame'
|
|
114
|
+
|
|
115
|
+
monkeypatch.setattr(pl, 'read_parquet', fake_read_parquet)
|
|
116
|
+
|
|
117
|
+
assert s3.get_dataset('s3://demo-bucket/data/').to_polars() == 'frame'
|
|
118
|
+
|
|
119
|
+
provider = captured['credential_provider']
|
|
120
|
+
assert provider is not None, 'Polars was left to find its own credentials'
|
|
121
|
+
values, _expiry = provider()
|
|
122
|
+
assert values == {
|
|
123
|
+
'aws_access_key_id': 'AKIA-first',
|
|
124
|
+
'aws_secret_access_key': 'secret-first',
|
|
125
|
+
'aws_session_token': 'token-first',
|
|
126
|
+
}
|
|
127
|
+
assert captured['storage_options'] == {'aws_region': 'us-east-1'}
|
|
128
|
+
|
|
129
|
+
|
|
130
|
+
def test_the_polars_provider_re_reads_credentials_on_every_call(loaded, monkeypatch):
|
|
131
|
+
"""A provider that captured its keys once would break when they rotate."""
|
|
132
|
+
import polars as pl
|
|
133
|
+
|
|
134
|
+
captured: dict = {}
|
|
135
|
+
monkeypatch.setattr(pl, 'scan_parquet', lambda source, **kwargs: captured.update(kwargs) or 'lazy')
|
|
136
|
+
|
|
137
|
+
_session, credentials, _connection = loaded
|
|
138
|
+
s3.get_dataset('s3://demo-bucket/data/').to_polars(lazy=True)
|
|
139
|
+
provider = captured['credential_provider']
|
|
140
|
+
|
|
141
|
+
assert provider()[0]['aws_access_key_id'] == 'AKIA-first'
|
|
142
|
+
credentials.rotate()
|
|
143
|
+
assert provider()[0]['aws_access_key_id'] == 'AKIA-second'
|
|
144
|
+
|
|
145
|
+
|
|
146
|
+
def test_from_polars_passes_credentials_too(loaded, monkeypatch):
|
|
147
|
+
"""Also guards the partitioning API: it is built for real here.
|
|
148
|
+
|
|
149
|
+
`PartitionMaxSize` was removed in the Polars release this package requires,
|
|
150
|
+
which broke this method outright — an AttributeError before any credential
|
|
151
|
+
was ever consulted.
|
|
152
|
+
"""
|
|
153
|
+
import polars as pl
|
|
154
|
+
|
|
155
|
+
captured: dict = {}
|
|
156
|
+
monkeypatch.setattr(
|
|
157
|
+
pl.DataFrame, 'write_parquet', lambda self, target, **kwargs: captured.update(target=target, **kwargs)
|
|
158
|
+
)
|
|
159
|
+
monkeypatch.setattr(s3.S3Dataset, 'clear_contents', lambda self: None)
|
|
160
|
+
|
|
161
|
+
s3.get_dataset('s3://demo-bucket/data/').from_polars(pl.DataFrame({'a': [1]}))
|
|
162
|
+
|
|
163
|
+
assert isinstance(captured['target'], pl.PartitionBy)
|
|
164
|
+
assert captured['credential_provider'] is not None
|
|
165
|
+
assert captured['storage_options'] == {'aws_region': 'us-east-1'}
|
|
166
|
+
|
|
167
|
+
|
|
168
|
+
# -- pyarrow ---------------------------------------------------------------
|
|
169
|
+
def test_to_arrow_reads_through_the_component_filesystem(loaded, monkeypatch):
|
|
170
|
+
"""pyarrow would otherwise build a filesystem from its own credential chain."""
|
|
171
|
+
import pyarrow.parquet as pq
|
|
172
|
+
|
|
173
|
+
captured: dict = {}
|
|
174
|
+
|
|
175
|
+
def fake_read_table(source, **kwargs):
|
|
176
|
+
captured.update(source=source, **kwargs)
|
|
177
|
+
return pa.table({'a': [1]})
|
|
178
|
+
|
|
179
|
+
monkeypatch.setattr(pq, 'read_table', fake_read_table)
|
|
180
|
+
|
|
181
|
+
session, _credentials, _connection = loaded
|
|
182
|
+
s3.get_dataset('s3://demo-bucket/data/').to_arrow()
|
|
183
|
+
|
|
184
|
+
assert captured['filesystem'] is session.get('s3').filesystem
|
|
185
|
+
# pyarrow does not strip the scheme when given a filesystem: passing the URI
|
|
186
|
+
# would address a key literally named "s3:".
|
|
187
|
+
assert captured['source'] == 'demo-bucket/data/'
|
|
188
|
+
|
|
189
|
+
|
|
190
|
+
def test_from_arrow_writes_through_the_component_filesystem(loaded, monkeypatch):
|
|
191
|
+
import pyarrow.parquet as pq
|
|
192
|
+
|
|
193
|
+
captured: dict = {}
|
|
194
|
+
monkeypatch.setattr(
|
|
195
|
+
pq, 'write_to_dataset', lambda table, root_path, **kwargs: captured.update(root_path=root_path, **kwargs)
|
|
196
|
+
)
|
|
197
|
+
monkeypatch.setattr(s3.S3Dataset, 'clear_contents', lambda self: None)
|
|
198
|
+
|
|
199
|
+
session, _credentials, _connection = loaded
|
|
200
|
+
s3.get_dataset('s3://demo-bucket/data/').from_arrow(pa.table({'a': [1]}))
|
|
201
|
+
|
|
202
|
+
assert captured['filesystem'] is session.get('s3').filesystem
|
|
203
|
+
assert captured['root_path'] == 'demo-bucket/data/'
|
|
204
|
+
|
|
205
|
+
|
|
206
|
+
def test_the_filesystem_is_rebuilt_when_credentials_rotate(loaded):
|
|
207
|
+
"""pyarrow copies the keys at construction and never refreshes them."""
|
|
208
|
+
session, credentials, _connection = loaded
|
|
209
|
+
component = session.get('s3')
|
|
210
|
+
|
|
211
|
+
first = component.filesystem
|
|
212
|
+
assert component.filesystem is first, 'unchanged credentials should not rebuild'
|
|
213
|
+
|
|
214
|
+
credentials.rotate()
|
|
215
|
+
assert component.filesystem is not first
|
|
216
|
+
|
|
217
|
+
|
|
218
|
+
# -- DuckDB ----------------------------------------------------------------
|
|
219
|
+
def test_the_duckdb_secret_follows_rotated_credentials(loaded):
|
|
220
|
+
"""Explicit keys are a snapshot inside DuckDB; they have to be refreshed."""
|
|
221
|
+
_session, credentials, connection = loaded
|
|
222
|
+
before = len(connection.secrets())
|
|
223
|
+
|
|
224
|
+
s3.get_dataset('s3://demo-bucket/data/').to_duckdb()
|
|
225
|
+
assert len(connection.secrets()) == before, 'unchanged credentials should not re-register'
|
|
226
|
+
|
|
227
|
+
credentials.rotate()
|
|
228
|
+
s3.get_dataset('s3://demo-bucket/data/').to_duckdb()
|
|
229
|
+
assert connection.secrets()[-1].count('AKIA-second') == 1
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|