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.
- apache_airflow_providers_datasluice-0.1.2/.gitignore +72 -0
- apache_airflow_providers_datasluice-0.1.2/PKG-INFO +19 -0
- apache_airflow_providers_datasluice-0.1.2/pyproject.toml +59 -0
- apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/__init__.py +1 -0
- apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/get_provider_info.py +47 -0
- apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/hooks/__init__.py +5 -0
- apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/hooks/datasluice.py +377 -0
- apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/operators/__init__.py +6 -0
- apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/operators/_xcom.py +26 -0
- apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/operators/materialize.py +62 -0
- apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/operators/search.py +59 -0
- apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/provider.yaml +21 -0
- apache_airflow_providers_datasluice-0.1.2/tests/dags/example_datasluice.py +26 -0
- apache_airflow_providers_datasluice-0.1.2/tests/measure_coverage.py +96 -0
- apache_airflow_providers_datasluice-0.1.2/tests/run_candidate.py +299 -0
- apache_airflow_providers_datasluice-0.1.2/tests/smoke.py +34 -0
- apache_airflow_providers_datasluice-0.1.2/tests/test_hook.py +366 -0
- apache_airflow_providers_datasluice-0.1.2/tests/test_materialize_operator.py +205 -0
- apache_airflow_providers_datasluice-0.1.2/tests/test_provider_metadata.py +154 -0
- apache_airflow_providers_datasluice-0.1.2/tests/test_public_boundary.py +44 -0
- 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."""
|
apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/get_provider_info.py
ADDED
|
@@ -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
|
+
}
|
apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/hooks/datasluice.py
ADDED
|
@@ -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
|
apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/operators/__init__.py
ADDED
|
@@ -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
|
apache_airflow_providers_datasluice-0.1.2/src/airflow/providers/datasluice/operators/materialize.py
ADDED
|
@@ -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
|