pytest-databricks-lakeflow 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- pytest_databricks_lakeflow-0.1.0/.gitignore +11 -0
- pytest_databricks_lakeflow-0.1.0/PKG-INFO +248 -0
- pytest_databricks_lakeflow-0.1.0/README.md +227 -0
- pytest_databricks_lakeflow-0.1.0/examples/bundle/databricks.yml +22 -0
- pytest_databricks_lakeflow-0.1.0/examples/bundle/pyproject.toml +15 -0
- pytest_databricks_lakeflow-0.1.0/examples/bundle/resources/my_example_job.job.yml +43 -0
- pytest_databricks_lakeflow-0.1.0/examples/bundle/src/lakeflow_example/__init__.py +0 -0
- pytest_databricks_lakeflow-0.1.0/examples/bundle/src/lakeflow_example/main.py +6 -0
- pytest_databricks_lakeflow-0.1.0/examples/bundle/src/notebooks/extract.py +4 -0
- pytest_databricks_lakeflow-0.1.0/examples/bundle/src/notebooks/load.py +2 -0
- pytest_databricks_lakeflow-0.1.0/examples/bundle/src/notebooks/transform.py +28 -0
- pytest_databricks_lakeflow-0.1.0/pyproject.toml +68 -0
- pytest_databricks_lakeflow-0.1.0/src/pytest_databricks_lakeflow/__init__.py +20 -0
- pytest_databricks_lakeflow-0.1.0/src/pytest_databricks_lakeflow/bundle.py +215 -0
- pytest_databricks_lakeflow-0.1.0/src/pytest_databricks_lakeflow/client.py +228 -0
- pytest_databricks_lakeflow-0.1.0/src/pytest_databricks_lakeflow/ephemeral.py +215 -0
- pytest_databricks_lakeflow-0.1.0/src/pytest_databricks_lakeflow/errors.py +13 -0
- pytest_databricks_lakeflow-0.1.0/src/pytest_databricks_lakeflow/helpers.py +59 -0
- pytest_databricks_lakeflow-0.1.0/src/pytest_databricks_lakeflow/hooks.py +30 -0
- pytest_databricks_lakeflow-0.1.0/src/pytest_databricks_lakeflow/options.py +122 -0
- pytest_databricks_lakeflow-0.1.0/src/pytest_databricks_lakeflow/plugin.py +140 -0
- pytest_databricks_lakeflow-0.1.0/src/pytest_databricks_lakeflow/py.typed +0 -0
- pytest_databricks_lakeflow-0.1.0/src/pytest_databricks_lakeflow/results.py +374 -0
- pytest_databricks_lakeflow-0.1.0/src/pytest_databricks_lakeflow/session.py +167 -0
- pytest_databricks_lakeflow-0.1.0/src/pytest_databricks_lakeflow/settings.py +341 -0
- pytest_databricks_lakeflow-0.1.0/tests/conftest.py +3 -0
- pytest_databricks_lakeflow-0.1.0/tests/integration/__init__.py +0 -0
- pytest_databricks_lakeflow-0.1.0/tests/integration/test_acceptance.py +101 -0
- pytest_databricks_lakeflow-0.1.0/tests/unit/__init__.py +0 -0
- pytest_databricks_lakeflow-0.1.0/tests/unit/conftest.py +59 -0
- pytest_databricks_lakeflow-0.1.0/tests/unit/data/runs/cancelled.json +192 -0
- pytest_databricks_lakeflow-0.1.0/tests/unit/data/runs/failed_no_probe_leaf.json +602 -0
- pytest_databricks_lakeflow-0.1.0/tests/unit/data/runs/only_plus_transform.json +427 -0
- pytest_databricks_lakeflow-0.1.0/tests/unit/data/runs/probe_failure.json +480 -0
- pytest_databricks_lakeflow-0.1.0/tests/unit/data/runs/probe_success.json +442 -0
- pytest_databricks_lakeflow-0.1.0/tests/unit/data/runs/probe_upstream_failed.json +735 -0
- pytest_databricks_lakeflow-0.1.0/tests/unit/data/runs/wheel_logs.json +1382 -0
- pytest_databricks_lakeflow-0.1.0/tests/unit/data/task_errors.json +177 -0
- pytest_databricks_lakeflow-0.1.0/tests/unit/fakes.py +234 -0
- pytest_databricks_lakeflow-0.1.0/tests/unit/test_build_job_run.py +153 -0
- pytest_databricks_lakeflow-0.1.0/tests/unit/test_bundle.py +214 -0
- pytest_databricks_lakeflow-0.1.0/tests/unit/test_cleanup.py +203 -0
- pytest_databricks_lakeflow-0.1.0/tests/unit/test_ephemeral.py +385 -0
- pytest_databricks_lakeflow-0.1.0/tests/unit/test_plugin_options.py +94 -0
- pytest_databricks_lakeflow-0.1.0/tests/unit/test_results.py +130 -0
- pytest_databricks_lakeflow-0.1.0/tests/unit/test_session.py +274 -0
- pytest_databricks_lakeflow-0.1.0/tests/unit/test_settings.py +155 -0
- pytest_databricks_lakeflow-0.1.0/tests/unit/test_task_values.py +163 -0
|
@@ -0,0 +1,248 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: pytest-databricks-lakeflow
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Run Lakeflow Jobs from a Databricks Asset Bundle as throwaway jobs inside pytest tests.
|
|
5
|
+
Project-URL: Homepage, https://github.com/stevenayers/pytest-databricks-lakeflow
|
|
6
|
+
Project-URL: Issues, https://github.com/stevenayers/pytest-databricks-lakeflow/issues
|
|
7
|
+
Author-email: Steven Ayers <steven@ayers.io>
|
|
8
|
+
Classifier: Framework :: Pytest
|
|
9
|
+
Classifier: Programming Language :: Python :: 3
|
|
10
|
+
Classifier: Programming Language :: Python :: 3 :: Only
|
|
11
|
+
Classifier: Topic :: Software Development :: Testing
|
|
12
|
+
Requires-Python: >=3.11
|
|
13
|
+
Requires-Dist: databricks-sdk>=0.54
|
|
14
|
+
Requires-Dist: filelock>=3
|
|
15
|
+
Requires-Dist: pytest>=8
|
|
16
|
+
Provides-Extra: pytester
|
|
17
|
+
Requires-Dist: databricks-labs-pytester>=0.7; extra == 'pytester'
|
|
18
|
+
Provides-Extra: xdist
|
|
19
|
+
Requires-Dist: pytest-xdist>=3; extra == 'xdist'
|
|
20
|
+
Description-Content-Type: text/markdown
|
|
21
|
+
|
|
22
|
+
# pytest-databricks-lakeflow
|
|
23
|
+
|
|
24
|
+
Write ordinary pytest tests that run a Lakeflow Job, or a single task from one, defined in your Databricks Asset Bundle. Each test gets its own throwaway copy of the job, runs it with the parameters you choose, and asserts on the outcome: success or failure, task outputs, logs, task values, and the tables the job wrote.
|
|
25
|
+
|
|
26
|
+
```python
|
|
27
|
+
import pytest
|
|
28
|
+
|
|
29
|
+
from pytest_databricks_lakeflow import Lakeflow
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def test_job_succeeds(lakeflow: Lakeflow) -> None:
|
|
33
|
+
with lakeflow.job("my_example_job", params={"param1": "hello", "param2": "world"}) as job:
|
|
34
|
+
run = job.run()
|
|
35
|
+
assert run.succeeded
|
|
36
|
+
assert run.task("transform").output == "hello world"
|
|
37
|
+
assert "Success!" in run.task("my_wheel_task").logs
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def test_job_surfaces_permission_error(lakeflow: Lakeflow) -> None:
|
|
41
|
+
with lakeflow.job("my_example_job", params={"param1": "forbidden"}) as job:
|
|
42
|
+
run = job.run()
|
|
43
|
+
assert run.failed
|
|
44
|
+
assert run.termination_code == "RUN_EXECUTION_ERROR"
|
|
45
|
+
assert run.task("transform").raised(PermissionError)
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
def test_subset_of_job(lakeflow: Lakeflow) -> None:
|
|
49
|
+
with lakeflow.job("my_example_job") as job:
|
|
50
|
+
run = job.run(only=["+transform"])
|
|
51
|
+
assert run.succeeded
|
|
52
|
+
assert "load" not in run.tasks
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
@pytest.mark.parametrize(("greeting", "expected"), [("hello", "hello world"), ("goodbye", "goodbye world")])
|
|
56
|
+
def test_single_task(lakeflow: Lakeflow, lakeflow_schema: str, greeting: str, expected: str) -> None:
|
|
57
|
+
with lakeflow.task(
|
|
58
|
+
"my_example_job.transform",
|
|
59
|
+
params={"param1": greeting, "param2": "world", "schema": lakeflow_schema},
|
|
60
|
+
upstream_values={"extract": {"row_count": 42, "source_path": "/Volumes/x/y/z"}},
|
|
61
|
+
capture_values=["rows_written"],
|
|
62
|
+
) as task:
|
|
63
|
+
run = task.run()
|
|
64
|
+
assert run.succeeded
|
|
65
|
+
assert run.task("transform").output == expected
|
|
66
|
+
assert run.task("transform").values["rows_written"] == 42
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def test_run_twice_in_one_job(lakeflow: Lakeflow, lakeflow_schema: str) -> None:
|
|
70
|
+
with lakeflow.job("my_example_job", params={"schema": lakeflow_schema}) as job:
|
|
71
|
+
first = job.run()
|
|
72
|
+
second = job.run(params={"mode": "incremental"})
|
|
73
|
+
assert first.succeeded and second.succeeded
|
|
74
|
+
assert second.duration < first.duration
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
def test_values_captured_even_on_failure(lakeflow: Lakeflow) -> None:
|
|
78
|
+
with lakeflow.job("my_example_job", params={"param1": "explode"}, capture_values={"transform": ["stage"]}) as job:
|
|
79
|
+
run = job.run()
|
|
80
|
+
assert run.failed
|
|
81
|
+
assert run.task("transform").values["stage"] == "validation"
|
|
82
|
+
```
|
|
83
|
+
|
|
84
|
+
[`examples/bundle`](https://github.com/stevenayers/pytest-databricks-lakeflow/tree/main/examples/bundle) is a complete bundle that these tests run against.
|
|
85
|
+
|
|
86
|
+
## Install
|
|
87
|
+
|
|
88
|
+
```bash
|
|
89
|
+
pip install pytest-databricks-lakeflow
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
You also need:
|
|
93
|
+
|
|
94
|
+
- Python 3.11 or later, pytest 8 or later and `databricks-sdk` 0.54 or later. These are installed as dependencies.
|
|
95
|
+
- The [Databricks CLI](https://docs.databricks.com/dev-tools/cli/install) v1.7.0 or later on `PATH`. The plugin shells out to it for `bundle deploy`, `summary`, `validate`, `sync` and `destroy`.
|
|
96
|
+
- Credentials that both the CLI and the SDK can use: a profile in `~/.databrickscfg` or the standard `DATABRICKS_*` environment variables.
|
|
97
|
+
- Serverless jobs in the workspace for the helper tasks, or see `lakeflow_helpers_compute` below.
|
|
98
|
+
|
|
99
|
+
For parallel tests, install the extra: `pip install "pytest-databricks-lakeflow[xdist]"`.
|
|
100
|
+
|
|
101
|
+
## Set up
|
|
102
|
+
|
|
103
|
+
Add a target for tests to `databricks.yml`. Development mode pauses schedules and triggers and lets runs overlap, and the separate `root_path` keeps the test deployment away from everyone else's:
|
|
104
|
+
|
|
105
|
+
```yaml
|
|
106
|
+
targets:
|
|
107
|
+
test:
|
|
108
|
+
mode: development
|
|
109
|
+
workspace:
|
|
110
|
+
root_path: ~/.bundle/${bundle.name}/test
|
|
111
|
+
presets:
|
|
112
|
+
trigger_pause_status: PAUSED
|
|
113
|
+
tags:
|
|
114
|
+
purpose: pytest
|
|
115
|
+
```
|
|
116
|
+
|
|
117
|
+
Configure the plugin in `pytest.ini`, `pyproject.toml` or `tox.ini`:
|
|
118
|
+
|
|
119
|
+
```ini
|
|
120
|
+
[pytest]
|
|
121
|
+
lakeflow_target = test
|
|
122
|
+
lakeflow_catalog = dev_sandbox
|
|
123
|
+
lakeflow_timeout_minutes = 20
|
|
124
|
+
```
|
|
125
|
+
|
|
126
|
+
The job code under test must read its catalog and schema from **job parameters**, not from bundle variables. Bundle variables are fixed at deploy time, so passing a different value per test only works through parameters.
|
|
127
|
+
|
|
128
|
+
## How it works
|
|
129
|
+
|
|
130
|
+
1. The first test that uses `lakeflow` deploys the bundle once with `databricks bundle deploy -t test`, reads the deployed job ids from `databricks bundle summary`, and loads each job's settings with `jobs.get`. Those settings have resolved workspace paths for notebooks and uploaded wheels. It also uploads two small helper notebooks to `/Users/<you>/.pytest-databricks-lakeflow/`.
|
|
131
|
+
2. `lakeflow.job(key)` copies the settings of the bundle job `key` (the resource key, as used by `databricks bundle run`). It removes the schedule, triggers, notifications, `run_as` and the bundle deployment marker, allows one run at a time, names the job `[pytest] <test node id>` and tags it `pytest`, `pytest_run_id` and `pytest_nodeid`.
|
|
132
|
+
3. Entering the `with` block creates the job, and `run()` runs it and waits. You can call `run()` more than once.
|
|
133
|
+
4. Leaving the block deletes the job. If the block exited with an exception, other than `pytest.skip` or `pytest.xfail`, the job is kept so you can open the run in the UI. The job and run URLs are logged and added to the test report.
|
|
134
|
+
|
|
135
|
+
Your tests never run the deployed jobs directly.
|
|
136
|
+
|
|
137
|
+
`lakeflow.job()` and `lakeflow.task()` also work without `with`. The job is then created by the first `run()` and deleted at the end of the test, or kept if the test failed.
|
|
138
|
+
|
|
139
|
+
### Single tasks and task values
|
|
140
|
+
|
|
141
|
+
`lakeflow.task("job_key.task_key", ...)` builds a job that holds only that task. For each upstream task named in `upstream_values` it adds a stub task with the same key, which sets those task values. Both `dbutils.jobs.taskValues.get(taskKey="extract", ...)` and `{{tasks.extract.values.row_count}}` in the task's settings then resolve against the stub.
|
|
142
|
+
|
|
143
|
+
Task values cannot be read through the Jobs API, so `capture_values` adds a probe task that runs after the target, even when it fails, and reports the values. `lakeflow.job(key, capture_values={"transform": ["stage"]})` does the same for a whole job.
|
|
144
|
+
|
|
145
|
+
When a value is missing, `run.task("t").values["k"]` raises a `KeyError` that says why: the task never set it, the task did not run, or the probe failed.
|
|
146
|
+
|
|
147
|
+
Stub values must be JSON-serialisable, at most 49,152 characters of JSON each, with keys of at most 512 characters. The plugin checks this before creating the job. If the target task refers to task values of an upstream task you did not stub, you get an error before anything is created.
|
|
148
|
+
|
|
149
|
+
## API
|
|
150
|
+
|
|
151
|
+
### Fixtures
|
|
152
|
+
|
|
153
|
+
- `lakeflow`: a `Lakeflow` for the current test.
|
|
154
|
+
- `lakeflow_schema`: creates `<lakeflow_catalog>.pytest_<run id>_<random>` and yields its full name. It is dropped after the test, even if the test failed, unless you pass `--lakeflow-keep`.
|
|
155
|
+
|
|
156
|
+
Tests that use either fixture get the `lakeflow` marker automatically, so `pytest -m "not lakeflow"` runs only your fast tests. No deploy happens in a session that runs no `lakeflow` test.
|
|
157
|
+
|
|
158
|
+
### `Lakeflow`
|
|
159
|
+
|
|
160
|
+
- `job(key, params=None, capture_values=None, name=None) -> EphemeralJob`. `params` are job parameters for every run. `capture_values` maps task keys to lists of value keys. `name` replaces the node id in the job name.
|
|
161
|
+
- `task("job_key.task_key", params=None, upstream_values=None, capture_values=None) -> EphemeralJob`. `upstream_values` maps upstream task keys to the values their stubs set. `capture_values` is a list of value keys.
|
|
162
|
+
- `template(key) -> JobSettings`: a fresh copy of the deployed settings.
|
|
163
|
+
- `ws`: the `WorkspaceClient`, for example to check that a table exists.
|
|
164
|
+
|
|
165
|
+
Parameters that the job does not declare raise an error, because the API would accept them silently. Set `lakeflow_allow_undeclared_params = true` to pass them anyway. Names starting with `__pytest_` are reserved.
|
|
166
|
+
|
|
167
|
+
### `EphemeralJob`
|
|
168
|
+
|
|
169
|
+
- `run(params=None, only=None, timeout=None) -> JobRun`. `params` are merged over the job's params for this run. `only` selects tasks, and `+task`, `task+` and `+task+` add upstream and downstream tasks. `timeout` is a `timedelta` that defaults to `lakeflow_timeout_minutes`. On timeout the run is cancelled and `LakeflowTimeout` is raised.
|
|
170
|
+
- `job_id`, `url` and `runs`.
|
|
171
|
+
|
|
172
|
+
### `JobRun`
|
|
173
|
+
|
|
174
|
+
- `succeeded`: true only if every task in `tasks` ended `SUCCESS`, apart from tasks `EXCLUDED` because their `run_if` or if/else condition was not met. `failed` is `not succeeded`.
|
|
175
|
+
- `tasks`: your tasks that were part of the run, by key. `skipped_tasks`: tasks left out by `only`. `injected_tasks`: the plugin's stubs and probes.
|
|
176
|
+
- `task(key)`: raises a `KeyError` that lists the available keys.
|
|
177
|
+
- `result_state`, `termination_code`, `termination_message`, `duration`, `run_id` and `run_url` are those of the parent run. Do not use the parent's `result_state` as the verdict: a job with a probe that runs after a failure ends `SUCCESS_WITH_FAILURES`.
|
|
178
|
+
|
|
179
|
+
### `TaskRun`
|
|
180
|
+
|
|
181
|
+
- `succeeded`, `result_state`, `termination_code`, `termination_message`, `duration`, `run_url` and `attempt`.
|
|
182
|
+
- `output`: the value passed to `dbutils.notebook.exit()`.
|
|
183
|
+
- `logs` and `log_lines`: stdout and stderr of Python script and wheel tasks.
|
|
184
|
+
- `error`: a `TaskError` with `type` (for example `"PermissionError"`), `message` and `trace`. `raised(PermissionError)` compares class names.
|
|
185
|
+
- `values` and `missing_values`: captured task values.
|
|
186
|
+
|
|
187
|
+
### Hooks
|
|
188
|
+
|
|
189
|
+
Implement these in a `conftest.py`. Like pytest's own test hooks, they apply to the tests in that conftest's directory and below:
|
|
190
|
+
|
|
191
|
+
- `pytest_lakeflow_modify_settings(settings, request)`: called just before each job is created. Change `settings` in place.
|
|
192
|
+
- `pytest_lakeflow_run_finished(run, request)`: called after each run.
|
|
193
|
+
|
|
194
|
+
## Configuration
|
|
195
|
+
|
|
196
|
+
| Ini option | Command line | Default | Meaning |
|
|
197
|
+
| --- | --- | --- | --- |
|
|
198
|
+
| `lakeflow_bundle_root` | | nearest directory with `databricks.yml` | Bundle root, relative to the rootdir. |
|
|
199
|
+
| `lakeflow_target` | `--lakeflow-target` | `test` | Bundle target. |
|
|
200
|
+
| `lakeflow_profile` | | standard SDK resolution | Profile for both the CLI and the SDK. |
|
|
201
|
+
| `lakeflow_sync` | `--lakeflow-sync` | `deploy` | `deploy`, `sync` or `none`. See below. |
|
|
202
|
+
| `lakeflow_timeout_minutes` | `--lakeflow-timeout` | `30` | Default timeout for one run. |
|
|
203
|
+
| `lakeflow_catalog` | | | Catalog for `lakeflow_schema`. |
|
|
204
|
+
| `lakeflow_helpers_path` | | `/Users/<you>/.pytest-databricks-lakeflow` | Workspace folder for the helper notebooks. |
|
|
205
|
+
| `lakeflow_helpers_compute` | | serverless | A `job_cluster_key` from the job for stub and probe tasks, for workspaces without serverless. |
|
|
206
|
+
| `lakeflow_helpers_environment_version` | | `4` | Serverless environment version for stub and probe tasks. |
|
|
207
|
+
| `lakeflow_disable_auto_optimization` | | `true` | Stop serverless from retrying failed tasks, so a failure fails once and fast. |
|
|
208
|
+
| `lakeflow_performance_optimized` | | `false` | Run every ephemeral job as `PERFORMANCE_OPTIMIZED`. |
|
|
209
|
+
| `lakeflow_allow_undeclared_params` | | `false` | Allow parameters the job does not declare. |
|
|
210
|
+
| `lakeflow_sweep_hours` | | `24` | Age after which `--lakeflow-sweep` deletes leftover jobs. |
|
|
211
|
+
| | `--lakeflow-keep` | off | Keep every job and schema. |
|
|
212
|
+
| | `--lakeflow-keep-failed`, `--no-lakeflow-keep-failed` | on | Keep the job when its test fails. |
|
|
213
|
+
| | `--lakeflow-sweep` | off | At session start, delete your own jobs tagged `pytest` from other runs that are older than `lakeflow_sweep_hours`. |
|
|
214
|
+
| | `--lakeflow-destroy` | off | Run `databricks bundle destroy` at the end of the session. |
|
|
215
|
+
|
|
216
|
+
The run id that tags every job comes from `LAKEFLOW_RUN_ID`, then `GITHUB_RUN_ID`, then `BUILD_BUILDID`, and is otherwise random.
|
|
217
|
+
|
|
218
|
+
### Sync strategies
|
|
219
|
+
|
|
220
|
+
- `deploy` (the default) deploys once per session and reads templates from the deployed jobs. This is the only strategy that builds and uploads wheels.
|
|
221
|
+
- `sync` runs `databricks bundle sync` and reads templates from `databricks bundle validate`. Jobs that use local wheels or JARs are refused, because `sync` does not upload them.
|
|
222
|
+
- `none` does nothing and reads templates from `validate`. Local wheel paths are pointed at the copies that an earlier `deploy` uploaded, after checking that they exist.
|
|
223
|
+
|
|
224
|
+
With `sync` and `none`, a job that uses `${resources...}` references is refused, because `validate` does not resolve them.
|
|
225
|
+
|
|
226
|
+
## Running in parallel
|
|
227
|
+
|
|
228
|
+
With pytest-xdist, one worker deploys and loads the templates while the others wait, then all of them reuse the result. Each worker uses its own `WorkspaceClient`.
|
|
229
|
+
|
|
230
|
+
With `databricks-cli` (U2M) authentication, several processes refreshing the token at once can fail with `exit status 45`. The plugin retries, but OAuth machine-to-machine credentials or a token are more reliable in CI.
|
|
231
|
+
|
|
232
|
+
## Timing
|
|
233
|
+
|
|
234
|
+
Measured in one Azure workspace:
|
|
235
|
+
|
|
236
|
+
- A trivial serverless notebook test takes about 20 to 35 seconds end to end, of which about 2 seconds is API overhead.
|
|
237
|
+
- A job with `performance_target: STANDARD` adds 3 to 4 minutes to every run. The plugin warns about it, and `lakeflow_performance_optimized = true` overrides it.
|
|
238
|
+
- A classic job cluster takes 2 to 10 minutes to start.
|
|
239
|
+
- Five jobs at once took no longer than one.
|
|
240
|
+
|
|
241
|
+
## Limitations
|
|
242
|
+
|
|
243
|
+
- **Logs.** `TaskRun.logs` is only filled for Python script, wheel and JAR tasks, and holds only the last 5 MB. Notebook tasks have no logs; use `output` instead. Notebook output keeps only the first 5 MB. The Databricks documentation says that on serverless compute the logs of one task can include output from other tasks in the same run. In our tests each task's logs held only its own output, but this is not guaranteed.
|
|
244
|
+
- **Task values** are not available through the Jobs API, so they are read by probe tasks. Each probe costs 7 to 30 seconds of serverless compute, even when its task did not run.
|
|
245
|
+
- **Job code must read its catalog and schema from job parameters**, as described above.
|
|
246
|
+
- **Single-task jobs** keep only the target task. On a shared job cluster, a task can import libraries that other tasks declared. On its own it loses them, so declare every library on the task that needs it.
|
|
247
|
+
- **Serverless auto-retry** is turned off in ephemeral jobs by default, so a test sees one attempt. Set `lakeflow_disable_auto_optimization = false` to keep the platform's behaviour. When a task is retried, `TaskRun` describes the last attempt.
|
|
248
|
+
- **Standard (shared) access mode** classic clusters were not tested. There, task values set from wheel or script tasks may be silently lost.
|
|
@@ -0,0 +1,227 @@
|
|
|
1
|
+
# pytest-databricks-lakeflow
|
|
2
|
+
|
|
3
|
+
Write ordinary pytest tests that run a Lakeflow Job, or a single task from one, defined in your Databricks Asset Bundle. Each test gets its own throwaway copy of the job, runs it with the parameters you choose, and asserts on the outcome: success or failure, task outputs, logs, task values, and the tables the job wrote.
|
|
4
|
+
|
|
5
|
+
```python
|
|
6
|
+
import pytest
|
|
7
|
+
|
|
8
|
+
from pytest_databricks_lakeflow import Lakeflow
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def test_job_succeeds(lakeflow: Lakeflow) -> None:
|
|
12
|
+
with lakeflow.job("my_example_job", params={"param1": "hello", "param2": "world"}) as job:
|
|
13
|
+
run = job.run()
|
|
14
|
+
assert run.succeeded
|
|
15
|
+
assert run.task("transform").output == "hello world"
|
|
16
|
+
assert "Success!" in run.task("my_wheel_task").logs
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def test_job_surfaces_permission_error(lakeflow: Lakeflow) -> None:
|
|
20
|
+
with lakeflow.job("my_example_job", params={"param1": "forbidden"}) as job:
|
|
21
|
+
run = job.run()
|
|
22
|
+
assert run.failed
|
|
23
|
+
assert run.termination_code == "RUN_EXECUTION_ERROR"
|
|
24
|
+
assert run.task("transform").raised(PermissionError)
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def test_subset_of_job(lakeflow: Lakeflow) -> None:
|
|
28
|
+
with lakeflow.job("my_example_job") as job:
|
|
29
|
+
run = job.run(only=["+transform"])
|
|
30
|
+
assert run.succeeded
|
|
31
|
+
assert "load" not in run.tasks
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
@pytest.mark.parametrize(("greeting", "expected"), [("hello", "hello world"), ("goodbye", "goodbye world")])
|
|
35
|
+
def test_single_task(lakeflow: Lakeflow, lakeflow_schema: str, greeting: str, expected: str) -> None:
|
|
36
|
+
with lakeflow.task(
|
|
37
|
+
"my_example_job.transform",
|
|
38
|
+
params={"param1": greeting, "param2": "world", "schema": lakeflow_schema},
|
|
39
|
+
upstream_values={"extract": {"row_count": 42, "source_path": "/Volumes/x/y/z"}},
|
|
40
|
+
capture_values=["rows_written"],
|
|
41
|
+
) as task:
|
|
42
|
+
run = task.run()
|
|
43
|
+
assert run.succeeded
|
|
44
|
+
assert run.task("transform").output == expected
|
|
45
|
+
assert run.task("transform").values["rows_written"] == 42
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
def test_run_twice_in_one_job(lakeflow: Lakeflow, lakeflow_schema: str) -> None:
|
|
49
|
+
with lakeflow.job("my_example_job", params={"schema": lakeflow_schema}) as job:
|
|
50
|
+
first = job.run()
|
|
51
|
+
second = job.run(params={"mode": "incremental"})
|
|
52
|
+
assert first.succeeded and second.succeeded
|
|
53
|
+
assert second.duration < first.duration
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
def test_values_captured_even_on_failure(lakeflow: Lakeflow) -> None:
|
|
57
|
+
with lakeflow.job("my_example_job", params={"param1": "explode"}, capture_values={"transform": ["stage"]}) as job:
|
|
58
|
+
run = job.run()
|
|
59
|
+
assert run.failed
|
|
60
|
+
assert run.task("transform").values["stage"] == "validation"
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
[`examples/bundle`](https://github.com/stevenayers/pytest-databricks-lakeflow/tree/main/examples/bundle) is a complete bundle that these tests run against.
|
|
64
|
+
|
|
65
|
+
## Install
|
|
66
|
+
|
|
67
|
+
```bash
|
|
68
|
+
pip install pytest-databricks-lakeflow
|
|
69
|
+
```
|
|
70
|
+
|
|
71
|
+
You also need:
|
|
72
|
+
|
|
73
|
+
- Python 3.11 or later, pytest 8 or later and `databricks-sdk` 0.54 or later. These are installed as dependencies.
|
|
74
|
+
- The [Databricks CLI](https://docs.databricks.com/dev-tools/cli/install) v1.7.0 or later on `PATH`. The plugin shells out to it for `bundle deploy`, `summary`, `validate`, `sync` and `destroy`.
|
|
75
|
+
- Credentials that both the CLI and the SDK can use: a profile in `~/.databrickscfg` or the standard `DATABRICKS_*` environment variables.
|
|
76
|
+
- Serverless jobs in the workspace for the helper tasks, or see `lakeflow_helpers_compute` below.
|
|
77
|
+
|
|
78
|
+
For parallel tests, install the extra: `pip install "pytest-databricks-lakeflow[xdist]"`.
|
|
79
|
+
|
|
80
|
+
## Set up
|
|
81
|
+
|
|
82
|
+
Add a target for tests to `databricks.yml`. Development mode pauses schedules and triggers and lets runs overlap, and the separate `root_path` keeps the test deployment away from everyone else's:
|
|
83
|
+
|
|
84
|
+
```yaml
|
|
85
|
+
targets:
|
|
86
|
+
test:
|
|
87
|
+
mode: development
|
|
88
|
+
workspace:
|
|
89
|
+
root_path: ~/.bundle/${bundle.name}/test
|
|
90
|
+
presets:
|
|
91
|
+
trigger_pause_status: PAUSED
|
|
92
|
+
tags:
|
|
93
|
+
purpose: pytest
|
|
94
|
+
```
|
|
95
|
+
|
|
96
|
+
Configure the plugin in `pytest.ini`, `pyproject.toml` or `tox.ini`:
|
|
97
|
+
|
|
98
|
+
```ini
|
|
99
|
+
[pytest]
|
|
100
|
+
lakeflow_target = test
|
|
101
|
+
lakeflow_catalog = dev_sandbox
|
|
102
|
+
lakeflow_timeout_minutes = 20
|
|
103
|
+
```
|
|
104
|
+
|
|
105
|
+
The job code under test must read its catalog and schema from **job parameters**, not from bundle variables. Bundle variables are fixed at deploy time, so passing a different value per test only works through parameters.
|
|
106
|
+
|
|
107
|
+
## How it works
|
|
108
|
+
|
|
109
|
+
1. The first test that uses `lakeflow` deploys the bundle once with `databricks bundle deploy -t test`, reads the deployed job ids from `databricks bundle summary`, and loads each job's settings with `jobs.get`. Those settings have resolved workspace paths for notebooks and uploaded wheels. It also uploads two small helper notebooks to `/Users/<you>/.pytest-databricks-lakeflow/`.
|
|
110
|
+
2. `lakeflow.job(key)` copies the settings of the bundle job `key` (the resource key, as used by `databricks bundle run`). It removes the schedule, triggers, notifications, `run_as` and the bundle deployment marker, allows one run at a time, names the job `[pytest] <test node id>` and tags it `pytest`, `pytest_run_id` and `pytest_nodeid`.
|
|
111
|
+
3. Entering the `with` block creates the job, and `run()` runs it and waits. You can call `run()` more than once.
|
|
112
|
+
4. Leaving the block deletes the job. If the block exited with an exception, other than `pytest.skip` or `pytest.xfail`, the job is kept so you can open the run in the UI. The job and run URLs are logged and added to the test report.
|
|
113
|
+
|
|
114
|
+
Your tests never run the deployed jobs directly.
|
|
115
|
+
|
|
116
|
+
`lakeflow.job()` and `lakeflow.task()` also work without `with`. The job is then created by the first `run()` and deleted at the end of the test, or kept if the test failed.
|
|
117
|
+
|
|
118
|
+
### Single tasks and task values
|
|
119
|
+
|
|
120
|
+
`lakeflow.task("job_key.task_key", ...)` builds a job that holds only that task. For each upstream task named in `upstream_values` it adds a stub task with the same key, which sets those task values. Both `dbutils.jobs.taskValues.get(taskKey="extract", ...)` and `{{tasks.extract.values.row_count}}` in the task's settings then resolve against the stub.
|
|
121
|
+
|
|
122
|
+
Task values cannot be read through the Jobs API, so `capture_values` adds a probe task that runs after the target, even when it fails, and reports the values. `lakeflow.job(key, capture_values={"transform": ["stage"]})` does the same for a whole job.
|
|
123
|
+
|
|
124
|
+
When a value is missing, `run.task("t").values["k"]` raises a `KeyError` that says why: the task never set it, the task did not run, or the probe failed.
|
|
125
|
+
|
|
126
|
+
Stub values must be JSON-serialisable, at most 49,152 characters of JSON each, with keys of at most 512 characters. The plugin checks this before creating the job. If the target task refers to task values of an upstream task you did not stub, you get an error before anything is created.
|
|
127
|
+
|
|
128
|
+
## API
|
|
129
|
+
|
|
130
|
+
### Fixtures
|
|
131
|
+
|
|
132
|
+
- `lakeflow`: a `Lakeflow` for the current test.
|
|
133
|
+
- `lakeflow_schema`: creates `<lakeflow_catalog>.pytest_<run id>_<random>` and yields its full name. It is dropped after the test, even if the test failed, unless you pass `--lakeflow-keep`.
|
|
134
|
+
|
|
135
|
+
Tests that use either fixture get the `lakeflow` marker automatically, so `pytest -m "not lakeflow"` runs only your fast tests. No deploy happens in a session that runs no `lakeflow` test.
|
|
136
|
+
|
|
137
|
+
### `Lakeflow`
|
|
138
|
+
|
|
139
|
+
- `job(key, params=None, capture_values=None, name=None) -> EphemeralJob`. `params` are job parameters for every run. `capture_values` maps task keys to lists of value keys. `name` replaces the node id in the job name.
|
|
140
|
+
- `task("job_key.task_key", params=None, upstream_values=None, capture_values=None) -> EphemeralJob`. `upstream_values` maps upstream task keys to the values their stubs set. `capture_values` is a list of value keys.
|
|
141
|
+
- `template(key) -> JobSettings`: a fresh copy of the deployed settings.
|
|
142
|
+
- `ws`: the `WorkspaceClient`, for example to check that a table exists.
|
|
143
|
+
|
|
144
|
+
Parameters that the job does not declare raise an error, because the API would accept them silently. Set `lakeflow_allow_undeclared_params = true` to pass them anyway. Names starting with `__pytest_` are reserved.
|
|
145
|
+
|
|
146
|
+
### `EphemeralJob`
|
|
147
|
+
|
|
148
|
+
- `run(params=None, only=None, timeout=None) -> JobRun`. `params` are merged over the job's params for this run. `only` selects tasks, and `+task`, `task+` and `+task+` add upstream and downstream tasks. `timeout` is a `timedelta` that defaults to `lakeflow_timeout_minutes`. On timeout the run is cancelled and `LakeflowTimeout` is raised.
|
|
149
|
+
- `job_id`, `url` and `runs`.
|
|
150
|
+
|
|
151
|
+
### `JobRun`
|
|
152
|
+
|
|
153
|
+
- `succeeded`: true only if every task in `tasks` ended `SUCCESS`, apart from tasks `EXCLUDED` because their `run_if` or if/else condition was not met. `failed` is `not succeeded`.
|
|
154
|
+
- `tasks`: your tasks that were part of the run, by key. `skipped_tasks`: tasks left out by `only`. `injected_tasks`: the plugin's stubs and probes.
|
|
155
|
+
- `task(key)`: raises a `KeyError` that lists the available keys.
|
|
156
|
+
- `result_state`, `termination_code`, `termination_message`, `duration`, `run_id` and `run_url` are those of the parent run. Do not use the parent's `result_state` as the verdict: a job with a probe that runs after a failure ends `SUCCESS_WITH_FAILURES`.
|
|
157
|
+
|
|
158
|
+
### `TaskRun`
|
|
159
|
+
|
|
160
|
+
- `succeeded`, `result_state`, `termination_code`, `termination_message`, `duration`, `run_url` and `attempt`.
|
|
161
|
+
- `output`: the value passed to `dbutils.notebook.exit()`.
|
|
162
|
+
- `logs` and `log_lines`: stdout and stderr of Python script and wheel tasks.
|
|
163
|
+
- `error`: a `TaskError` with `type` (for example `"PermissionError"`), `message` and `trace`. `raised(PermissionError)` compares class names.
|
|
164
|
+
- `values` and `missing_values`: captured task values.
|
|
165
|
+
|
|
166
|
+
### Hooks
|
|
167
|
+
|
|
168
|
+
Implement these in a `conftest.py`. Like pytest's own test hooks, they apply to the tests in that conftest's directory and below:
|
|
169
|
+
|
|
170
|
+
- `pytest_lakeflow_modify_settings(settings, request)`: called just before each job is created. Change `settings` in place.
|
|
171
|
+
- `pytest_lakeflow_run_finished(run, request)`: called after each run.
|
|
172
|
+
|
|
173
|
+
## Configuration
|
|
174
|
+
|
|
175
|
+
| Ini option | Command line | Default | Meaning |
|
|
176
|
+
| --- | --- | --- | --- |
|
|
177
|
+
| `lakeflow_bundle_root` | | nearest directory with `databricks.yml` | Bundle root, relative to the rootdir. |
|
|
178
|
+
| `lakeflow_target` | `--lakeflow-target` | `test` | Bundle target. |
|
|
179
|
+
| `lakeflow_profile` | | standard SDK resolution | Profile for both the CLI and the SDK. |
|
|
180
|
+
| `lakeflow_sync` | `--lakeflow-sync` | `deploy` | `deploy`, `sync` or `none`. See below. |
|
|
181
|
+
| `lakeflow_timeout_minutes` | `--lakeflow-timeout` | `30` | Default timeout for one run. |
|
|
182
|
+
| `lakeflow_catalog` | | | Catalog for `lakeflow_schema`. |
|
|
183
|
+
| `lakeflow_helpers_path` | | `/Users/<you>/.pytest-databricks-lakeflow` | Workspace folder for the helper notebooks. |
|
|
184
|
+
| `lakeflow_helpers_compute` | | serverless | A `job_cluster_key` from the job for stub and probe tasks, for workspaces without serverless. |
|
|
185
|
+
| `lakeflow_helpers_environment_version` | | `4` | Serverless environment version for stub and probe tasks. |
|
|
186
|
+
| `lakeflow_disable_auto_optimization` | | `true` | Stop serverless from retrying failed tasks, so a failure fails once and fast. |
|
|
187
|
+
| `lakeflow_performance_optimized` | | `false` | Run every ephemeral job as `PERFORMANCE_OPTIMIZED`. |
|
|
188
|
+
| `lakeflow_allow_undeclared_params` | | `false` | Allow parameters the job does not declare. |
|
|
189
|
+
| `lakeflow_sweep_hours` | | `24` | Age after which `--lakeflow-sweep` deletes leftover jobs. |
|
|
190
|
+
| | `--lakeflow-keep` | off | Keep every job and schema. |
|
|
191
|
+
| | `--lakeflow-keep-failed`, `--no-lakeflow-keep-failed` | on | Keep the job when its test fails. |
|
|
192
|
+
| | `--lakeflow-sweep` | off | At session start, delete your own jobs tagged `pytest` from other runs that are older than `lakeflow_sweep_hours`. |
|
|
193
|
+
| | `--lakeflow-destroy` | off | Run `databricks bundle destroy` at the end of the session. |
|
|
194
|
+
|
|
195
|
+
The run id that tags every job comes from `LAKEFLOW_RUN_ID`, then `GITHUB_RUN_ID`, then `BUILD_BUILDID`, and is otherwise random.
|
|
196
|
+
|
|
197
|
+
### Sync strategies
|
|
198
|
+
|
|
199
|
+
- `deploy` (the default) deploys once per session and reads templates from the deployed jobs. This is the only strategy that builds and uploads wheels.
|
|
200
|
+
- `sync` runs `databricks bundle sync` and reads templates from `databricks bundle validate`. Jobs that use local wheels or JARs are refused, because `sync` does not upload them.
|
|
201
|
+
- `none` does nothing and reads templates from `validate`. Local wheel paths are pointed at the copies that an earlier `deploy` uploaded, after checking that they exist.
|
|
202
|
+
|
|
203
|
+
With `sync` and `none`, a job that uses `${resources...}` references is refused, because `validate` does not resolve them.
|
|
204
|
+
|
|
205
|
+
## Running in parallel
|
|
206
|
+
|
|
207
|
+
With pytest-xdist, one worker deploys and loads the templates while the others wait, then all of them reuse the result. Each worker uses its own `WorkspaceClient`.
|
|
208
|
+
|
|
209
|
+
With `databricks-cli` (U2M) authentication, several processes refreshing the token at once can fail with `exit status 45`. The plugin retries, but OAuth machine-to-machine credentials or a token are more reliable in CI.
|
|
210
|
+
|
|
211
|
+
## Timing
|
|
212
|
+
|
|
213
|
+
Measured in one Azure workspace:
|
|
214
|
+
|
|
215
|
+
- A trivial serverless notebook test takes about 20 to 35 seconds end to end, of which about 2 seconds is API overhead.
|
|
216
|
+
- A job with `performance_target: STANDARD` adds 3 to 4 minutes to every run. The plugin warns about it, and `lakeflow_performance_optimized = true` overrides it.
|
|
217
|
+
- A classic job cluster takes 2 to 10 minutes to start.
|
|
218
|
+
- Five jobs at once took no longer than one.
|
|
219
|
+
|
|
220
|
+
## Limitations
|
|
221
|
+
|
|
222
|
+
- **Logs.** `TaskRun.logs` is only filled for Python script, wheel and JAR tasks, and holds only the last 5 MB. Notebook tasks have no logs; use `output` instead. Notebook output keeps only the first 5 MB. The Databricks documentation says that on serverless compute the logs of one task can include output from other tasks in the same run. In our tests each task's logs held only its own output, but this is not guaranteed.
|
|
223
|
+
- **Task values** are not available through the Jobs API, so they are read by probe tasks. Each probe costs 7 to 30 seconds of serverless compute, even when its task did not run.
|
|
224
|
+
- **Job code must read its catalog and schema from job parameters**, as described above.
|
|
225
|
+
- **Single-task jobs** keep only the target task. On a shared job cluster, a task can import libraries that other tasks declared. On its own it loses them, so declare every library on the task that needs it.
|
|
226
|
+
- **Serverless auto-retry** is turned off in ephemeral jobs by default, so a test sees one attempt. Set `lakeflow_disable_auto_optimization = false` to keep the platform's behaviour. When a task is retried, `TaskRun` describes the last attempt.
|
|
227
|
+
- **Standard (shared) access mode** classic clusters were not tested. There, task values set from wheel or script tasks may be silently lost.
|
|
@@ -0,0 +1,22 @@
|
|
|
1
|
+
bundle:
|
|
2
|
+
name: lakeflow_example
|
|
3
|
+
|
|
4
|
+
include:
|
|
5
|
+
- resources/*.yml
|
|
6
|
+
|
|
7
|
+
artifacts:
|
|
8
|
+
lakeflow_example:
|
|
9
|
+
type: whl
|
|
10
|
+
build: uv build --wheel
|
|
11
|
+
path: .
|
|
12
|
+
|
|
13
|
+
targets:
|
|
14
|
+
test:
|
|
15
|
+
mode: development
|
|
16
|
+
default: true
|
|
17
|
+
workspace:
|
|
18
|
+
root_path: ~/.bundle/${bundle.name}/test
|
|
19
|
+
presets:
|
|
20
|
+
trigger_pause_status: PAUSED
|
|
21
|
+
tags:
|
|
22
|
+
purpose: pytest
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "lakeflow_example"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "Wheel task for the pytest-databricks-lakeflow example bundle."
|
|
5
|
+
requires-python = ">=3.10"
|
|
6
|
+
|
|
7
|
+
[project.scripts]
|
|
8
|
+
main = "lakeflow_example.main:main"
|
|
9
|
+
|
|
10
|
+
[build-system]
|
|
11
|
+
requires = ["hatchling"]
|
|
12
|
+
build-backend = "hatchling.build"
|
|
13
|
+
|
|
14
|
+
[tool.hatch.build.targets.wheel]
|
|
15
|
+
packages = ["src/lakeflow_example"]
|
|
@@ -0,0 +1,43 @@
|
|
|
1
|
+
resources:
|
|
2
|
+
jobs:
|
|
3
|
+
my_example_job:
|
|
4
|
+
name: my_example_job
|
|
5
|
+
description: Example job for the pytest-databricks-lakeflow integration tests.
|
|
6
|
+
parameters:
|
|
7
|
+
- name: param1
|
|
8
|
+
default: hello
|
|
9
|
+
- name: param2
|
|
10
|
+
default: world
|
|
11
|
+
- name: schema
|
|
12
|
+
default: ""
|
|
13
|
+
- name: mode
|
|
14
|
+
default: full
|
|
15
|
+
environments:
|
|
16
|
+
- environment_key: default
|
|
17
|
+
spec:
|
|
18
|
+
environment_version: "4"
|
|
19
|
+
dependencies:
|
|
20
|
+
- ../dist/*.whl
|
|
21
|
+
tasks:
|
|
22
|
+
- task_key: extract
|
|
23
|
+
notebook_task:
|
|
24
|
+
notebook_path: ../src/notebooks/extract.py
|
|
25
|
+
- task_key: transform
|
|
26
|
+
depends_on:
|
|
27
|
+
- task_key: extract
|
|
28
|
+
notebook_task:
|
|
29
|
+
notebook_path: ../src/notebooks/transform.py
|
|
30
|
+
base_parameters:
|
|
31
|
+
row_count: "{{tasks.extract.values.row_count}}"
|
|
32
|
+
- task_key: load
|
|
33
|
+
depends_on:
|
|
34
|
+
- task_key: transform
|
|
35
|
+
notebook_task:
|
|
36
|
+
notebook_path: ../src/notebooks/load.py
|
|
37
|
+
- task_key: my_wheel_task
|
|
38
|
+
environment_key: default
|
|
39
|
+
python_wheel_task:
|
|
40
|
+
package_name: lakeflow_example
|
|
41
|
+
entry_point: main
|
|
42
|
+
parameters:
|
|
43
|
+
- "{{job.parameters.param1}}"
|
|
File without changes
|
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
# Databricks notebook source
|
|
2
|
+
import time
|
|
3
|
+
|
|
4
|
+
# Job parameters arrive as widgets. Catalog and schema must come from job parameters, not bundle variables,
|
|
5
|
+
# so that each test can point the job at its own temporary schema.
|
|
6
|
+
param1 = dbutils.widgets.get("param1")
|
|
7
|
+
param2 = dbutils.widgets.get("param2")
|
|
8
|
+
schema = dbutils.widgets.get("schema")
|
|
9
|
+
mode = dbutils.widgets.get("mode")
|
|
10
|
+
row_count = int(dbutils.widgets.get("row_count"))
|
|
11
|
+
source_path = dbutils.jobs.taskValues.get(taskKey="extract", key="source_path")
|
|
12
|
+
|
|
13
|
+
dbutils.jobs.taskValues.set(key="stage", value="validation")
|
|
14
|
+
dbutils.jobs.taskValues.set(key="source_path", value=source_path)
|
|
15
|
+
if param1 == "forbidden":
|
|
16
|
+
raise PermissionError(f"{param1} may not read {source_path}")
|
|
17
|
+
if param1 == "explode":
|
|
18
|
+
raise ValueError(f"validation failed for {param1}")
|
|
19
|
+
|
|
20
|
+
if schema:
|
|
21
|
+
spark.range(row_count).write.mode("overwrite").saveAsTable(f"{schema}.example_rows")
|
|
22
|
+
if mode == "full":
|
|
23
|
+
# A full load is deliberately slow so that the example tests can compare run durations.
|
|
24
|
+
time.sleep(30)
|
|
25
|
+
|
|
26
|
+
dbutils.jobs.taskValues.set(key="rows_written", value=row_count)
|
|
27
|
+
dbutils.jobs.taskValues.set(key="stage", value="done")
|
|
28
|
+
dbutils.notebook.exit(f"{param1} {param2}")
|