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.
Files changed (46) hide show
  1. {arrows-0.2.0 → arrows-0.2.2}/PKG-INFO +7 -2
  2. {arrows-0.2.0 → arrows-0.2.2}/README.md +5 -1
  3. {arrows-0.2.0 → arrows-0.2.2}/pyproject.toml +2 -1
  4. {arrows-0.2.0 → arrows-0.2.2}/pyproject.toml.orig +2 -1
  5. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/components/aws.py +69 -6
  6. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/components/s3.py +28 -17
  7. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/s3.py +52 -12
  8. arrows-0.2.2/tests/test_s3_credentials.py +229 -0
  9. {arrows-0.2.0 → arrows-0.2.2}/LICENSE +0 -0
  10. {arrows-0.2.0 → arrows-0.2.2}/examples/README.md +0 -0
  11. {arrows-0.2.0 → arrows-0.2.2}/examples/bootstrap_aws.py +0 -0
  12. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/__init__.py +0 -0
  13. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/__main__.py +0 -0
  14. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/auth.py +0 -0
  15. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/cli.py +0 -0
  16. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/components/__init__.py +0 -0
  17. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/components/gmail.py +0 -0
  18. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/components/google.py +0 -0
  19. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/components/google_sheets.py +0 -0
  20. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/components/redshift.py +0 -0
  21. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/components/sqlite.py +0 -0
  22. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/core/__init__.py +0 -0
  23. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/core/component.py +0 -0
  24. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/core/config.py +0 -0
  25. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/core/engine.py +0 -0
  26. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/core/errors.py +0 -0
  27. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/core/registry.py +0 -0
  28. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/core/secrets.py +0 -0
  29. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/core/session.py +0 -0
  30. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/gmail.py +0 -0
  31. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/google_sheets.py +0 -0
  32. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/mysql.py +0 -0
  33. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/py.typed +0 -0
  34. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/redis.py +0 -0
  35. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/redshift.py +0 -0
  36. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/spark.py +0 -0
  37. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/sqlite.py +0 -0
  38. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/template_renderer.py +0 -0
  39. {arrows-0.2.0 → arrows-0.2.2}/src/arrows/utils.py +0 -0
  40. {arrows-0.2.0 → arrows-0.2.2}/tests/conftest.py +0 -0
  41. {arrows-0.2.0 → arrows-0.2.2}/tests/test_login.py +0 -0
  42. {arrows-0.2.0 → arrows-0.2.2}/tests/test_packaging.py +0 -0
  43. {arrows-0.2.0 → arrows-0.2.2}/tests/test_public_api.py +0 -0
  44. {arrows-0.2.0 → arrows-0.2.2}/tests/test_registry_and_session.py +0 -0
  45. {arrows-0.2.0 → arrows-0.2.2}/tests/test_secrets.py +0 -0
  46. {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.0
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 src tests # lint
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 src tests # lint
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.0"
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.0"
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. They are then visible in ``duckdb_secrets()``, which is the
118
- cost of that setup and a reason to prefer a real AWS profile.
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
- credentials = self.boto3_session.get_credentials()
125
- if credentials is None:
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
- frozen = credentials.get_frozen_credentials()
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
- """Lazily built :class:`pyarrow.fs.S3FileSystem` bound to the AWS session."""
30
- if self._filesystem is None:
31
- from pyarrow.fs import S3FileSystem
32
-
33
- aws = self.session.get('aws')
34
- credentials = aws.boto3_session.get_credentials()
35
- frozen = credentials.get_frozen_credentials() if credentials else None
36
- kwargs: dict[str, Any] = {}
37
- if frozen is not None:
38
- kwargs = {
39
- 'access_key': frozen.access_key,
40
- 'secret_key': frozen.secret_key,
41
- 'session_token': frozen.token,
42
- }
43
- if aws.region:
44
- kwargs['region'] = aws.region
45
- self._filesystem = S3FileSystem(**kwargs)
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.s3_path)
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 get_duckdb().from_parquet(f'{self.s3_path}*.parquet')
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
- return pl.scan_parquet(self.s3_path) if lazy else pl.read_parquet(self.s3_path)
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.s3_path)
153
+ pq.write_to_dataset(arrow, self._key, filesystem=self.s3)
118
154
  else:
119
- get_duckdb().execute(f"COPY arrow TO '{self.s3_path[:-1]}' (FORMAT parquet, FILE_SIZE_BYTES '1G')")
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
- partition_info = pl.PartitionMaxSize(base_path=self.s3_path, max_size=512_000)
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 get_duckdb().sql(f'WITH temp_s3_dataset_table AS (SELECT * FROM {source})\n{body}') # noqa: S608
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 get_duckdb().sql(rendered)
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.s3_path[5:]
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.s3_path[5:]
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