apache-airflow-providers-datasluice 0.1.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 (21) hide show
  1. apache_airflow_providers_datasluice-0.1.2/.gitignore +72 -0
  2. apache_airflow_providers_datasluice-0.1.2/PKG-INFO +19 -0
  3. apache_airflow_providers_datasluice-0.1.2/pyproject.toml +59 -0
  4. apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/__init__.py +1 -0
  5. apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/get_provider_info.py +47 -0
  6. apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/hooks/__init__.py +5 -0
  7. apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/hooks/datasluice.py +377 -0
  8. apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/operators/__init__.py +6 -0
  9. apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/operators/_xcom.py +26 -0
  10. apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/operators/materialize.py +62 -0
  11. apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/operators/search.py +59 -0
  12. apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/provider.yaml +21 -0
  13. apache_airflow_providers_datasluice-0.1.2/tests/dags/example_datasluice.py +26 -0
  14. apache_airflow_providers_datasluice-0.1.2/tests/measure_coverage.py +96 -0
  15. apache_airflow_providers_datasluice-0.1.2/tests/run_candidate.py +299 -0
  16. apache_airflow_providers_datasluice-0.1.2/tests/smoke.py +34 -0
  17. apache_airflow_providers_datasluice-0.1.2/tests/test_hook.py +366 -0
  18. apache_airflow_providers_datasluice-0.1.2/tests/test_materialize_operator.py +205 -0
  19. apache_airflow_providers_datasluice-0.1.2/tests/test_provider_metadata.py +154 -0
  20. apache_airflow_providers_datasluice-0.1.2/tests/test_public_boundary.py +44 -0
  21. apache_airflow_providers_datasluice-0.1.2/tests/test_search_operator.py +160 -0
@@ -0,0 +1,72 @@
1
+ # Byte-compiled / optimized / DLL files
2
+ __pycache__/
3
+ *.py[codz]
4
+ *$py.class
5
+
6
+ # C extensions
7
+ *.so
8
+
9
+ # Distribution / packaging
10
+ build/
11
+ dist/
12
+ eggs/
13
+ .eggs/
14
+ *.egg-info/
15
+ *.egg
16
+ MANIFEST
17
+
18
+ # Unit test / coverage reports
19
+ htmlcov/
20
+ .tox/
21
+ .nox/
22
+ .coverage
23
+ .coverage.*
24
+ coverage.xml
25
+ .hypothesis/
26
+ .pytest_cache/
27
+
28
+ # Environments
29
+ .env
30
+ .envrc
31
+ .venv
32
+ .opencode
33
+
34
+ # uv
35
+ # Generally recommended to include uv.lock in version control for
36
+ # reproducibility. Uncomment for libraries where you want consumers
37
+ # to resolve their own dependencies.
38
+ # uv.lock
39
+
40
+ # Ruff
41
+ .ruff_cache/
42
+
43
+ # Type checkers
44
+ .mypy_cache/
45
+ .dmypy.json
46
+ dmypy.json
47
+ .pyre/
48
+ .pytype/
49
+
50
+ # Documentation build output
51
+ site/
52
+
53
+ # IDE files
54
+ # .idea/
55
+ # .vscode/
56
+
57
+ # GSD planning
58
+ .planning/
59
+
60
+ # Data
61
+ /data/
62
+
63
+ # Notebooks
64
+ notebook/
65
+
66
+ # Graphify
67
+ graphify-out/*
68
+ !graphify-out/.graphify_labels.json
69
+ !graphify-out/GRAPH_REPORT.md
70
+ !graphify-out/graph.html
71
+ !graphify-out/graph.json
72
+ !graphify-out/manifest.json
@@ -0,0 +1,19 @@
1
+ Metadata-Version: 2.4
2
+ Name: apache-airflow-providers-datasluice
3
+ Version: 0.1.2
4
+ Summary: Apache Airflow provider for DataSluice open-data discovery, streaming, and materialization
5
+ Project-URL: bugs, https://github.com/bocklabs/datasluice/issues
6
+ Project-URL: changelog, https://github.com/bocklabs/datasluice/releases
7
+ Project-URL: documentation, https://bocklabs.github.io/datasluice/
8
+ Project-URL: homepage, https://github.com/bocklabs/datasluice
9
+ Author-email: Nitish Raj <contact@bocklabs.com>
10
+ Maintainer-email: Nitish Raj <contact@bocklabs.com>
11
+ License: MIT
12
+ Classifier: Framework :: Apache Airflow
13
+ Classifier: Framework :: Apache Airflow :: Provider
14
+ Classifier: License :: OSI Approved :: MIT License
15
+ Classifier: Programming Language :: Python :: 3
16
+ Classifier: Typing :: Typed
17
+ Requires-Python: >=3.12
18
+ Requires-Dist: apache-airflow<4,>=3.2
19
+ Requires-Dist: datasluice<1,>=0.2
@@ -0,0 +1,59 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "apache-airflow-providers-datasluice"
7
+ version = "0.1.2"
8
+ description = "Apache Airflow provider for DataSluice open-data discovery, streaming, and materialization"
9
+ requires-python = ">=3.12"
10
+ license = {text = "MIT"}
11
+ authors = [
12
+ {name = "Nitish Raj", email = "contact@bocklabs.com"}
13
+ ]
14
+ maintainers = [
15
+ {name = "Nitish Raj", email = "contact@bocklabs.com"}
16
+ ]
17
+ classifiers = [
18
+ "License :: OSI Approved :: MIT License",
19
+ "Programming Language :: Python :: 3",
20
+ "Framework :: Apache Airflow",
21
+ "Framework :: Apache Airflow :: Provider",
22
+ "Typing :: Typed",
23
+ ]
24
+ dependencies = [
25
+ "datasluice>=0.2,<1",
26
+ "apache-airflow>=3.2,<4",
27
+ ]
28
+
29
+ [project.entry-points."apache_airflow_provider"]
30
+ provider_info = "airflow.providers.datasluice.get_provider_info:get_provider_info"
31
+
32
+ [project.urls]
33
+ bugs = "https://github.com/bocklabs/datasluice/issues"
34
+ changelog = "https://github.com/bocklabs/datasluice/releases"
35
+ documentation = "https://bocklabs.github.io/datasluice/"
36
+ homepage = "https://github.com/bocklabs/datasluice"
37
+
38
+ [tool.hatch.build.targets.wheel]
39
+ sources = ["src"]
40
+ packages = ["src/airflow/providers/datasluice"]
41
+
42
+ [tool.hatch.build.targets.sdist]
43
+ include = [
44
+ "src/airflow/providers/datasluice",
45
+ "tests",
46
+ "pyproject.toml",
47
+ ]
48
+
49
+ [tool.coverage.run]
50
+ branch = true
51
+ parallel = true
52
+ source = ["airflow/providers/datasluice"]
53
+ data_file = ".coverage.provider"
54
+
55
+ [tool.coverage.report]
56
+ show_missing = true
57
+ skip_covered = true
58
+ precision = 2
59
+ fail_under = 80.00
@@ -0,0 +1 @@
1
+ """DataSluice Apache Airflow provider package."""
@@ -0,0 +1,47 @@
1
+ """Native Airflow provider discovery callable for DataSluice.
2
+
3
+ Returns declarative metadata consumed by Airflow's ``ProviderManager`` via the
4
+ ``apache_airflow_provider`` entry point. The callable is metadata-only: it performs
5
+ no Connection lookup, network I/O, or private-core import.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ import importlib.metadata
11
+ from typing import Any
12
+
13
+ _PROVIDER_PACKAGE = "apache-airflow-providers-datasluice"
14
+ _PROVIDER_NAME = "DataSluice"
15
+ _PROVIDER_DESCRIPTION = "Apache Airflow provider for DataSluice open-data discovery, streaming, and materialization."
16
+ _HOOK_CLASS = "airflow.providers.datasluice.hooks.datasluice.DataSluiceHook"
17
+ _OPERATOR_MODULES = (
18
+ "airflow.providers.datasluice.operators.search",
19
+ "airflow.providers.datasluice.operators.materialize",
20
+ )
21
+
22
+
23
+ def get_provider_info() -> dict[str, Any]:
24
+ """Return the provider metadata dict consumed by Airflow ``ProviderManager``.
25
+
26
+ Returns:
27
+ A mapping with the locked package identity, hook, and connection metadata.
28
+ """
29
+ return {
30
+ "package-name": _PROVIDER_PACKAGE,
31
+ "name": _PROVIDER_NAME,
32
+ "description": _PROVIDER_DESCRIPTION,
33
+ "versions": [importlib.metadata.version(_PROVIDER_PACKAGE)],
34
+ "operators": [
35
+ {
36
+ "integration-name": _PROVIDER_NAME,
37
+ "python-modules": list(_OPERATOR_MODULES),
38
+ }
39
+ ],
40
+ "hook-class-names": [_HOOK_CLASS],
41
+ "connection-types": [
42
+ {
43
+ "hook-class-name": _HOOK_CLASS,
44
+ "connection-type": "datasluice",
45
+ }
46
+ ],
47
+ }
@@ -0,0 +1,5 @@
1
+ """DataSluice provider hooks."""
2
+
3
+ from airflow.providers.datasluice.hooks.datasluice import DataSluiceHook
4
+
5
+ __all__ = ["DataSluiceHook"]
@@ -0,0 +1,377 @@
1
+ """Connection-aware DataSluice Hook for Airflow 3."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import base64
6
+ from collections.abc import Mapping
7
+ from dataclasses import dataclass
8
+ from typing import Any
9
+ from urllib.parse import parse_qsl, urlsplit
10
+
11
+ from airflow.sdk import BaseHook
12
+
13
+ from datasluice import (
14
+ Artifact,
15
+ CatalogResourceLocator,
16
+ DataSluice,
17
+ DataSluiceError,
18
+ DirectResourceLocator,
19
+ ResourceLocator,
20
+ resource_locator_from_dict,
21
+ )
22
+
23
+ _SESSION_OPTIONS = frozenset({"timeout", "retries", "rate_limit", "page_size", "cache_dir", "cache_ttl"})
24
+ _OVERRIDE_OPTIONS = _SESSION_OPTIONS | frozenset({"portal_url", "portal_type"})
25
+ _SENSITIVE_QUERY_KEYS = frozenset({"access_token", "api_key", "auth", "password", "secret", "signature", "token"})
26
+ _DEFAULT_API_KEY_HEADER = "X-Api-Key"
27
+
28
+
29
+ class _ConnectionAuth:
30
+ """Apply one connection-owned credential strategy without exposing its values."""
31
+
32
+ def __init__(
33
+ self,
34
+ auth_type: str,
35
+ *,
36
+ secret: str | None = None,
37
+ username: str | None = None,
38
+ password: str | None = None,
39
+ header_name: str = _DEFAULT_API_KEY_HEADER,
40
+ param_name: str = "api_key",
41
+ in_header: bool = True,
42
+ in_query: bool = False,
43
+ scheme: str = "Bearer",
44
+ headers: Mapping[str, str] | None = None,
45
+ ) -> None:
46
+ self._auth_type = auth_type
47
+ self._secret = secret
48
+ self._username = username
49
+ self._password = password
50
+ self._header_name = header_name
51
+ self._param_name = param_name
52
+ self._in_header = in_header
53
+ self._in_query = in_query
54
+ self._scheme = scheme
55
+ self._headers = dict(headers or {})
56
+
57
+ def apply(
58
+ self, headers: dict[str, str], params: dict[str, Any] | None = None
59
+ ) -> tuple[dict[str, str], dict[str, Any]]:
60
+ applied_headers = dict(headers)
61
+ applied_params = dict(params or {})
62
+ if self._auth_type == "api_key":
63
+ if self._in_header and self._secret is not None:
64
+ applied_headers[self._header_name] = self._secret
65
+ if self._in_query and self._secret is not None:
66
+ applied_params[self._param_name] = self._secret
67
+ elif self._auth_type == "bearer" and self._secret is not None:
68
+ applied_headers["Authorization"] = f"{self._scheme} {self._secret}"
69
+ elif self._auth_type == "basic" and self._username is not None and self._password is not None:
70
+ token = base64.b64encode(f"{self._username}:{self._password}".encode()).decode("ascii")
71
+ applied_headers["Authorization"] = f"Basic {token}"
72
+ elif self._auth_type == "headers":
73
+ applied_headers.update(self._headers)
74
+ return applied_headers, applied_params
75
+
76
+ def __repr__(self) -> str:
77
+ return f"<_ConnectionAuth type={self._auth_type!r}>"
78
+
79
+
80
+ @dataclass(frozen=True)
81
+ class _ResolvedConnection:
82
+ """Execution-time, sanitized view of an Airflow Connection."""
83
+
84
+ portal_url: str | None
85
+ portal_type: str | None
86
+ session_kwargs: Mapping[str, object]
87
+ secrets: tuple[str, ...]
88
+
89
+
90
+ class DataSluiceHook(BaseHook):
91
+ """Connect Airflow tasks to the public DataSluice facade."""
92
+
93
+ conn_name_attr = "airflow_conn_id"
94
+ default_conn_name = "datasluice_default"
95
+ conn_type = "datasluice"
96
+ hook_name = "DataSluice"
97
+
98
+ def __init__(self, airflow_conn_id: str = default_conn_name, **overrides: object) -> None:
99
+ super().__init__()
100
+ if not isinstance(airflow_conn_id, str) or not airflow_conn_id:
101
+ raise ValueError("airflow_conn_id must be a non-empty string")
102
+ unsupported = set(overrides) - _OVERRIDE_OPTIONS
103
+ if unsupported:
104
+ raise ValueError("Only non-secret DataSluice connection overrides are supported")
105
+ for name, value in overrides.items():
106
+ _validate_option(name, value)
107
+ self.airflow_conn_id = airflow_conn_id
108
+ self._overrides = dict(overrides)
109
+
110
+ def get_conn(self) -> DataSluice:
111
+ """Build and return a DataSluice facade from the configured Connection."""
112
+ config = self._resolve_connection()
113
+ try:
114
+ return DataSluice(**config.session_kwargs)
115
+ except Exception as exc:
116
+ _raise_redacted(exc, config.secrets)
117
+
118
+ def search(
119
+ self,
120
+ portal_url: str | None = None,
121
+ query: object = None,
122
+ *,
123
+ max_results: int = 50,
124
+ portal_type: str | None = None,
125
+ **query_kwargs: object,
126
+ ) -> list[dict[str, object]]:
127
+ """Search a portal and return bounded, JSON-safe catalog locators."""
128
+ config = self._resolve_connection()
129
+ try:
130
+ limit = _validate_limit(max_results)
131
+ selected_url = _select_portal_url(portal_url, config.portal_url)
132
+ selected_portal_type = portal_type if portal_type is not None else config.portal_type
133
+ search_kwargs = dict(query_kwargs)
134
+ search_kwargs["limit"] = limit
135
+ if selected_portal_type is not None:
136
+ search_kwargs["portal_type"] = selected_portal_type
137
+ with DataSluice(**config.session_kwargs) as facade:
138
+ result = facade.search(selected_url, query, **search_kwargs)
139
+ return _catalog_locators(result, selected_url, limit)
140
+ except Exception as exc:
141
+ _raise_redacted(exc, config.secrets)
142
+
143
+ def materialize(
144
+ self,
145
+ locator: ResourceLocator | dict[str, object],
146
+ destination_uri: str,
147
+ *,
148
+ mode: str = "parquet",
149
+ ) -> Artifact:
150
+ """Materialize exactly one locator and return its canonical Artifact."""
151
+ selected_locator = resource_locator_from_dict(locator) if isinstance(locator, dict) else locator
152
+ if not isinstance(selected_locator, (DirectResourceLocator, CatalogResourceLocator)):
153
+ raise DataSluiceError("materialize requires one ResourceLocator")
154
+ config = self._resolve_connection()
155
+ try:
156
+ with DataSluice(**config.session_kwargs) as facade:
157
+ artifact = facade.materialize(selected_locator, destination_uri, mode=mode)
158
+ if not isinstance(artifact, Artifact):
159
+ raise DataSluiceError("DataSluice materialize did not return an Artifact")
160
+ return artifact
161
+ except Exception as exc:
162
+ _raise_redacted(exc, config.secrets)
163
+
164
+ def _resolve_connection(self) -> _ResolvedConnection:
165
+ connection = self.get_connection(self.airflow_conn_id)
166
+ extras = _connection_extras(connection)
167
+ portal_url = _resolve_text_option("portal_url", self._overrides, extras)
168
+ if portal_url is None:
169
+ portal_url = _first_text(extras, "portal_url", "base_url", "url") or _text(
170
+ getattr(connection, "host", None)
171
+ )
172
+ portal_type = _resolve_text_option("portal_type", self._overrides, extras)
173
+ if portal_type is None:
174
+ portal_type = _first_text(extras, "portal_type", "connector")
175
+
176
+ session_kwargs: dict[str, object] = {}
177
+ for name in _SESSION_OPTIONS:
178
+ value = _resolve_value(name, self._overrides, extras)
179
+ if value is not None:
180
+ session_kwargs[name] = _coerce_option(name, value)
181
+ auth, secrets = _connection_auth(connection, extras, portal_url)
182
+ if auth is not None:
183
+ session_kwargs["auth"] = auth
184
+ return _ResolvedConnection(
185
+ portal_url=portal_url,
186
+ portal_type=portal_type,
187
+ session_kwargs=session_kwargs,
188
+ secrets=secrets,
189
+ )
190
+
191
+
192
+ def _connection_extras(connection: object) -> dict[str, object]:
193
+ value = getattr(connection, "extra_dejson", {})
194
+ if not isinstance(value, Mapping):
195
+ return {}
196
+ return {key: item for key, item in value.items() if isinstance(key, str)}
197
+
198
+
199
+ def _validate_option(name: str, value: object) -> None:
200
+ if value is None:
201
+ return
202
+ if name in {"portal_url", "portal_type", "cache_dir"} and not isinstance(value, str):
203
+ raise ValueError(f"Invalid non-secret DataSluice override: {name}")
204
+ if name in {"retries", "page_size", "cache_ttl"} and (type(value) is not int or value < 0):
205
+ raise ValueError(f"Invalid non-secret DataSluice override: {name}")
206
+ if name in {"timeout", "rate_limit"} and (
207
+ isinstance(value, bool) or not isinstance(value, (int, float)) or value < 0
208
+ ):
209
+ raise ValueError(f"Invalid non-secret DataSluice override: {name}")
210
+
211
+
212
+ def _coerce_option(name: str, value: object) -> object:
213
+ _validate_option(name, value)
214
+ return value
215
+
216
+
217
+ def _resolve_value(name: str, overrides: Mapping[str, object], extras: Mapping[str, object]) -> object:
218
+ if name in overrides and overrides[name] is not None:
219
+ return overrides[name]
220
+ return extras.get(name)
221
+
222
+
223
+ def _resolve_text_option(name: str, overrides: Mapping[str, object], extras: Mapping[str, object]) -> str | None:
224
+ value = _resolve_value(name, overrides, extras)
225
+ if value is None:
226
+ return None
227
+ if not isinstance(value, str):
228
+ raise DataSluiceError(f"Invalid DataSluice connection option: {name}")
229
+ return value
230
+
231
+
232
+ def _first_text(values: Mapping[str, object], *names: str) -> str | None:
233
+ for name in names:
234
+ value = _text(values.get(name))
235
+ if value is not None:
236
+ return value
237
+ return None
238
+
239
+
240
+ def _text(value: object) -> str | None:
241
+ return value if isinstance(value, str) and value else None
242
+
243
+
244
+ def _connection_auth(
245
+ connection: object,
246
+ extras: Mapping[str, object],
247
+ portal_url: str | None,
248
+ ) -> tuple[_ConnectionAuth | None, tuple[str, ...]]:
249
+ login = _text(getattr(connection, "login", None)) or _text(extras.get("username"))
250
+ password = _text(getattr(connection, "password", None)) or _text(extras.get("password"))
251
+ api_key = _text(extras.get("api_key")) or password
252
+ token = _text(extras.get("token")) or _text(extras.get("bearer_token")) or password
253
+ raw_headers = extras.get("headers", extras.get("custom_headers"))
254
+ headers = _headers(raw_headers)
255
+ auth_type = _text(extras.get("auth_type"))
256
+ if auth_type is None:
257
+ if headers:
258
+ auth_type = "headers"
259
+ elif login and password:
260
+ auth_type = "basic"
261
+ elif api_key:
262
+ auth_type = "api_key"
263
+ normalized_type = auth_type.lower().replace("-", "_") if auth_type else "none"
264
+ secrets = _secret_values(login, password, api_key, token, *(headers.values() if headers else ()))
265
+ if normalized_type in {"none", "no_auth"}:
266
+ return None, _add_url_secrets(secrets, portal_url)
267
+ if normalized_type == "api_key":
268
+ if api_key is None:
269
+ raise DataSluiceError("DataSluice api_key authentication requires a Connection secret")
270
+ header_name = _text(extras.get("api_key_header")) or _DEFAULT_API_KEY_HEADER
271
+ param_name = _text(extras.get("api_key_param")) or "api_key"
272
+ in_header = _boolean_option(extras.get("api_key_in_header"), True, "api_key_in_header")
273
+ in_query = _boolean_option(extras.get("api_key_in_query"), False, "api_key_in_query")
274
+ return (
275
+ _ConnectionAuth(
276
+ "api_key",
277
+ secret=api_key,
278
+ header_name=header_name,
279
+ param_name=param_name,
280
+ in_header=in_header,
281
+ in_query=in_query,
282
+ ),
283
+ _add_url_secrets(secrets, portal_url),
284
+ )
285
+ if normalized_type == "bearer":
286
+ if token is None:
287
+ raise DataSluiceError("DataSluice bearer authentication requires a Connection secret")
288
+ scheme = _text(extras.get("bearer_scheme")) or "Bearer"
289
+ return _ConnectionAuth("bearer", secret=token, scheme=scheme), _add_url_secrets(secrets, portal_url)
290
+ if normalized_type == "basic":
291
+ if login is None or password is None:
292
+ raise DataSluiceError("DataSluice basic authentication requires Connection login and password")
293
+ return _ConnectionAuth("basic", username=login, password=password), _add_url_secrets(secrets, portal_url)
294
+ if normalized_type == "headers":
295
+ if not headers:
296
+ raise DataSluiceError("DataSluice headers authentication requires Connection headers")
297
+ return _ConnectionAuth("headers", headers=headers), _add_url_secrets(secrets, portal_url)
298
+ raise DataSluiceError("Unsupported DataSluice Connection auth_type")
299
+
300
+
301
+ def _headers(value: object) -> dict[str, str]:
302
+ if value is None:
303
+ return {}
304
+ if not isinstance(value, Mapping) or not all(
305
+ isinstance(key, str) and isinstance(item, str) for key, item in value.items()
306
+ ):
307
+ raise DataSluiceError("DataSluice Connection headers must be a string mapping")
308
+ return dict(value)
309
+
310
+
311
+ def _boolean_option(value: object, default: bool, name: str) -> bool:
312
+ if value is None:
313
+ return default
314
+ if not isinstance(value, bool):
315
+ raise DataSluiceError(f"Invalid DataSluice Connection option: {name}")
316
+ return value
317
+
318
+
319
+ def _secret_values(*values: str | None) -> tuple[str, ...]:
320
+ return tuple(dict.fromkeys(value for value in values if value))
321
+
322
+
323
+ def _add_url_secrets(secrets: tuple[str, ...], portal_url: str | None) -> tuple[str, ...]:
324
+ if portal_url is None:
325
+ return secrets
326
+ try:
327
+ query_secrets = tuple(
328
+ value
329
+ for name, value in parse_qsl(urlsplit(portal_url).query, keep_blank_values=True)
330
+ if name.lower() in _SENSITIVE_QUERY_KEYS and value
331
+ )
332
+ except ValueError:
333
+ query_secrets = ()
334
+ return _secret_values(*secrets, *query_secrets)
335
+
336
+
337
+ def _select_portal_url(explicit: str | None, configured: str | None) -> str:
338
+ selected = explicit or configured
339
+ if not selected:
340
+ raise DataSluiceError("DataSluice Connection has no portal_url")
341
+ return selected
342
+
343
+
344
+ def _validate_limit(value: int) -> int:
345
+ if type(value) is not int or not 1 <= value <= 1000:
346
+ raise DataSluiceError("DataSluice search max_results must be between 1 and 1000")
347
+ return value
348
+
349
+
350
+ def _catalog_locators(result: object, portal_url: str, limit: int) -> list[dict[str, object]]:
351
+ locators: list[dict[str, object]] = []
352
+ for dataset in getattr(result, "datasets", ()):
353
+ dataset_id = getattr(dataset, "id", None)
354
+ for resource in getattr(dataset, "resources", ()):
355
+ locator = CatalogResourceLocator(
356
+ portal_url=portal_url,
357
+ dataset_id=dataset_id,
358
+ resource_id=getattr(resource, "id", None),
359
+ )
360
+ locators.append(locator.to_dict())
361
+ if len(locators) == limit:
362
+ return locators
363
+ return locators
364
+
365
+
366
+ def _raise_redacted(error: Exception, secrets: tuple[str, ...]) -> None:
367
+ message = str(error)
368
+ redacted = message
369
+ for secret in secrets:
370
+ redacted = redacted.replace(secret, "***")
371
+ if redacted == message:
372
+ raise error
373
+ try:
374
+ replacement = type(error)(redacted)
375
+ except Exception:
376
+ replacement = RuntimeError(redacted)
377
+ raise replacement from None
@@ -0,0 +1,6 @@
1
+ """Airflow operators for DataSluice."""
2
+
3
+ from airflow.providers.datasluice.operators.materialize import DataSluiceMaterializeOperator
4
+ from airflow.providers.datasluice.operators.search import DataSluiceSearchOperator
5
+
6
+ __all__ = ["DataSluiceMaterializeOperator", "DataSluiceSearchOperator"]
@@ -0,0 +1,26 @@
1
+ """Bounded JSON encoding for provider XCom values."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import json
6
+
7
+ from datasluice import DataSluiceError
8
+
9
+ MAX_XCOM_BYTES = 49_152
10
+
11
+
12
+ def validate_xcom_payload[T](payload: T) -> T:
13
+ """Validate one deterministic, compact JSON XCom payload."""
14
+ try:
15
+ encoded = json.dumps(
16
+ payload,
17
+ allow_nan=False,
18
+ ensure_ascii=False,
19
+ separators=(",", ":"),
20
+ sort_keys=True,
21
+ ).encode("utf-8")
22
+ except (TypeError, ValueError, UnicodeEncodeError) as exc:
23
+ raise DataSluiceError("XCom payload must be JSON serializable") from exc
24
+ if len(encoded) > MAX_XCOM_BYTES:
25
+ raise DataSluiceError(f"XCom payload exceeds the {MAX_XCOM_BYTES}-byte limit")
26
+ return payload
@@ -0,0 +1,62 @@
1
+ """One-resource materialization operator for Airflow 3."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from typing import Any
6
+
7
+ from airflow.providers.datasluice.hooks.datasluice import DataSluiceHook
8
+ from airflow.providers.datasluice.operators._xcom import validate_xcom_payload
9
+ from airflow.sdk import BaseOperator
10
+
11
+ from datasluice import (
12
+ Artifact,
13
+ CatalogResourceLocator,
14
+ DataSluiceError,
15
+ DirectResourceLocator,
16
+ resource_locator_from_dict,
17
+ )
18
+
19
+
20
+ class DataSluiceMaterializeOperator(BaseOperator):
21
+ """Materialize one locator and return one bounded Artifact dictionary."""
22
+
23
+ template_fields = ("locator", "destination_uri", "mode", "airflow_conn_id")
24
+
25
+ def __init__(
26
+ self,
27
+ *,
28
+ locator: object,
29
+ destination_uri: str,
30
+ mode: str = "parquet",
31
+ airflow_conn_id: str = "datasluice_default",
32
+ **kwargs: Any,
33
+ ) -> None:
34
+ super().__init__(**kwargs)
35
+ self.locator = locator
36
+ self.destination_uri = destination_uri
37
+ self.mode = mode
38
+ self.airflow_conn_id = airflow_conn_id
39
+
40
+ def execute(self, context: Any) -> dict[str, object]:
41
+ """Materialize one resource and validate its XCom representation."""
42
+ del context
43
+ locator = _validate_locator(self.locator)
44
+ if not isinstance(self.destination_uri, str) or not self.destination_uri:
45
+ raise DataSluiceError("DataSluice materialize destination_uri must be a non-empty string")
46
+ if not isinstance(self.mode, str) or self.mode not in {"parquet", "raw"}:
47
+ raise DataSluiceError("DataSluice materialize mode must be parquet or raw")
48
+
49
+ hook = DataSluiceHook(airflow_conn_id=self.airflow_conn_id)
50
+ artifact = hook.materialize(locator, self.destination_uri, mode=self.mode)
51
+ if not isinstance(artifact, Artifact):
52
+ raise DataSluiceError("DataSluice materialize must return one Artifact")
53
+ return validate_xcom_payload(artifact.to_dict())
54
+
55
+
56
+ def _validate_locator(value: object) -> DirectResourceLocator | CatalogResourceLocator:
57
+ if not isinstance(value, dict):
58
+ raise DataSluiceError("DataSluice materialize requires exactly one locator dictionary")
59
+ locator = resource_locator_from_dict(value)
60
+ if not isinstance(locator, (DirectResourceLocator, CatalogResourceLocator)):
61
+ raise DataSluiceError("DataSluice materialize requires one ResourceLocator")
62
+ return locator