-
Notifications
You must be signed in to change notification settings - Fork 1.6k
/
base.py
494 lines (419 loc) · 16.2 KB
/
base.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
import os
import threading
import time
import traceback
from abc import ABCMeta, abstractmethod
from contextlib import nullcontext
from datetime import datetime
from typing import Type, Union, Dict, Any, Optional
import dbt.exceptions
from dbt import tracking
from dbt.adapters.factory import get_adapter
from dbt.config import RuntimeConfig, Project
from dbt.config.profile import read_profile
from dbt.contracts.graph.manifest import Manifest
from dbt.contracts.results import (
NodeStatus,
RunResult,
collect_timing_info,
RunStatus,
RunningStatus,
)
from dbt.events.contextvars import get_node_info
from dbt.events.functions import fire_event
from dbt.events.types import (
LogDbtProjectError,
LogDbtProfileError,
CatchableExceptionOnRun,
InternalErrorOnRun,
GenericExceptionOnRun,
NodeConnectionReleaseError,
LogDebugStackTrace,
SkippingDetails,
LogSkipBecauseError,
NodeCompiling,
NodeExecuting,
)
from dbt.exceptions import (
NotImplementedError,
CompilationError,
DbtRuntimeError,
DbtInternalError,
)
from dbt.flags import get_flags
from dbt.graph import Graph
from dbt.logger import log_manager
from .printer import print_run_result_error
from dbt.task.contextvars import cv_project_root
class NoneConfig:
@classmethod
def from_args(cls, args):
return None
def read_profiles(profiles_dir=None):
"""This is only used for some error handling"""
if profiles_dir is None:
profiles_dir = get_flags().PROFILES_DIR
raw_profiles = read_profile(profiles_dir)
if raw_profiles is None:
profiles = {}
else:
profiles = {k: v for (k, v) in raw_profiles.items() if k != "config"}
return profiles
class BaseTask(metaclass=ABCMeta):
ConfigType: Union[Type[NoneConfig], Type[Project]] = NoneConfig
def __init__(self, args, config, project=None):
self.args = args
self.config = config
self.project = config if isinstance(config, Project) else project
if self.config:
cv_project_root.set(self.config.project_root)
@classmethod
def pre_init_hook(cls, args):
"""A hook called before the task is initialized."""
if args.log_format == "json":
log_manager.format_json()
else:
log_manager.format_text()
@classmethod
def set_log_format(cls):
if get_flags().LOG_FORMAT == "json":
log_manager.format_json()
else:
log_manager.format_text()
@classmethod
def from_args(cls, args, *pargs, **kwargs):
try:
# This is usually RuntimeConfig
config = cls.ConfigType.from_args(args)
except dbt.exceptions.DbtProjectError as exc:
fire_event(LogDbtProjectError(exc=str(exc)))
tracking.track_invalid_invocation(args=args, result_type=exc.result_type)
raise dbt.exceptions.DbtRuntimeError("Could not run dbt") from exc
except dbt.exceptions.DbtProfileError as exc:
all_profile_names = list(read_profiles(get_flags().PROFILES_DIR).keys())
fire_event(LogDbtProfileError(exc=str(exc), profiles=all_profile_names))
tracking.track_invalid_invocation(args=args, result_type=exc.result_type)
raise dbt.exceptions.DbtRuntimeError("Could not run dbt") from exc
return cls(args, config, *pargs, **kwargs)
@abstractmethod
def run(self):
raise dbt.exceptions.NotImplementedError("Not Implemented")
def interpret_results(self, results):
return True
def get_nearest_project_dir(project_dir: Optional[str]) -> str:
# If the user provides an explicit project directory, use that
# but don't look at parent directories.
if project_dir:
project_file = os.path.join(project_dir, "dbt_project.yml")
if os.path.exists(project_file):
return project_dir
else:
raise dbt.exceptions.DbtRuntimeError(
"fatal: Invalid --project-dir flag. Not a dbt project. "
"Missing dbt_project.yml file"
)
root_path = os.path.abspath(os.sep)
cwd = os.getcwd()
while cwd != root_path:
project_file = os.path.join(cwd, "dbt_project.yml")
if os.path.exists(project_file):
return cwd
cwd = os.path.dirname(cwd)
raise dbt.exceptions.DbtRuntimeError(
"fatal: Not a dbt project (or any of the parent directories). "
"Missing dbt_project.yml file"
)
def move_to_nearest_project_dir(project_dir: Optional[str]) -> str:
nearest_project_dir = get_nearest_project_dir(project_dir)
os.chdir(nearest_project_dir)
return nearest_project_dir
# TODO: look into deprecating this class in favor of several small functions that
# produce the same behavior. currently this class only contains manifest compilation,
# holding a manifest, and moving direcories.
class ConfiguredTask(BaseTask):
ConfigType = RuntimeConfig
def __init__(self, args, config, manifest: Optional[Manifest] = None):
super().__init__(args, config)
self.graph: Optional[Graph] = None
self.manifest = manifest
def compile_manifest(self):
if self.manifest is None:
raise DbtInternalError("compile_manifest called before manifest was loaded")
start_compile_manifest = time.perf_counter()
# we cannot get adapter in init since it will break rpc #5579
adapter = get_adapter(self.config)
compiler = adapter.get_compiler()
self.graph = compiler.compile(self.manifest)
compile_time = time.perf_counter() - start_compile_manifest
if dbt.tracking.active_user is not None:
dbt.tracking.track_runnable_timing({"graph_compilation_elapsed": compile_time})
@classmethod
def from_args(cls, args, *pargs, **kwargs):
move_to_nearest_project_dir(args.project_dir)
return super().from_args(args, *pargs, **kwargs)
class ExecutionContext:
"""During execution and error handling, dbt makes use of mutable state:
timing information and the newest (compiled vs executed) form of the node.
"""
def __init__(self, node):
self.timing = []
self.node = node
class BaseRunner(metaclass=ABCMeta):
def __init__(self, config, adapter, node, node_index, num_nodes):
self.config = config
self.adapter = adapter
self.node = node
self.node_index = node_index
self.num_nodes = num_nodes
self.skip = False
self.skip_cause: Optional[RunResult] = None
self.run_ephemeral_models = False
@abstractmethod
def compile(self, manifest: Manifest) -> Any:
pass
def get_result_status(self, result) -> Dict[str, str]:
if result.status == NodeStatus.Error:
return {"node_status": "error", "node_error": str(result.message)}
elif result.status == NodeStatus.Skipped:
return {"node_status": "skipped"}
elif result.status == NodeStatus.Fail:
return {"node_status": "failed"}
elif result.status == NodeStatus.Warn:
return {"node_status": "warn"}
else:
return {"node_status": "passed"}
def run_with_hooks(self, manifest):
if self.skip:
return self.on_skip()
# no before/after printing for ephemeral mdoels
if not self.node.is_ephemeral_model:
self.before_execute()
result = self.safe_run(manifest)
self.node.update_event_status(
node_status=result.status, finished_at=datetime.utcnow().isoformat()
)
if not self.node.is_ephemeral_model:
self.after_execute(result)
return result
def _build_run_result(
self,
node,
start_time,
status,
timing_info,
message,
agate_table=None,
adapter_response=None,
failures=None,
):
execution_time = time.time() - start_time
thread_id = threading.current_thread().name
if adapter_response is None:
adapter_response = {}
return RunResult(
status=status,
thread_id=thread_id,
execution_time=execution_time,
timing=timing_info,
message=message,
node=node,
agate_table=agate_table,
adapter_response=adapter_response,
failures=failures,
)
def error_result(self, node, message, start_time, timing_info):
return self._build_run_result(
node=node,
start_time=start_time,
status=RunStatus.Error,
timing_info=timing_info,
message=message,
)
def ephemeral_result(self, node, start_time, timing_info):
return self._build_run_result(
node=node,
start_time=start_time,
status=RunStatus.Success,
timing_info=timing_info,
message=None,
)
def from_run_result(self, result, start_time, timing_info):
return self._build_run_result(
node=result.node,
start_time=start_time,
status=result.status,
timing_info=timing_info,
message=result.message,
agate_table=result.agate_table,
adapter_response=result.adapter_response,
failures=result.failures,
)
def skip_result(self, node, message):
thread_id = threading.current_thread().name
return RunResult(
status=RunStatus.Skipped,
thread_id=thread_id,
execution_time=0,
timing=[],
message=message,
node=node,
adapter_response={},
failures=None,
)
def compile_and_execute(self, manifest, ctx):
result = None
with self.adapter.connection_for(self.node) if get_flags().INTROSPECT else nullcontext():
ctx.node.update_event_status(node_status=RunningStatus.Compiling)
fire_event(
NodeCompiling(
node_info=ctx.node.node_info,
)
)
with collect_timing_info("compile", ctx.timing.append):
# if we fail here, we still have a compiled node to return
# this has the benefit of showing a build path for the errant
# model
ctx.node = self.compile(manifest)
# for ephemeral nodes, we only want to compile, not run
if not ctx.node.is_ephemeral_model or self.run_ephemeral_models:
ctx.node.update_event_status(node_status=RunningStatus.Executing)
fire_event(
NodeExecuting(
node_info=ctx.node.node_info,
)
)
with collect_timing_info("execute", ctx.timing.append):
result = self.run(ctx.node, manifest)
ctx.node = result.node
return result
def _handle_catchable_exception(self, e, ctx):
if e.node is None:
e.add_node(ctx.node)
fire_event(
CatchableExceptionOnRun(
exc=str(e), exc_info=traceback.format_exc(), node_info=get_node_info()
)
)
return str(e)
def _handle_internal_exception(self, e, ctx):
fire_event(InternalErrorOnRun(build_path=self.node.build_path, exc=str(e)))
return str(e)
def _handle_generic_exception(self, e, ctx):
fire_event(
GenericExceptionOnRun(
build_path=self.node.build_path,
unique_id=self.node.unique_id,
exc=str(e),
)
)
fire_event(LogDebugStackTrace(exc_info=traceback.format_exc()))
return str(e)
def handle_exception(self, e, ctx):
catchable_errors = (CompilationError, DbtRuntimeError)
if isinstance(e, catchable_errors):
error = self._handle_catchable_exception(e, ctx)
elif isinstance(e, DbtInternalError):
error = self._handle_internal_exception(e, ctx)
else:
error = self._handle_generic_exception(e, ctx)
return error
def safe_run(self, manifest):
started = time.time()
ctx = ExecutionContext(self.node)
error = None
result = None
try:
result = self.compile_and_execute(manifest, ctx)
except Exception as e:
error = self.handle_exception(e, ctx)
finally:
exc_str = self._safe_release_connection()
# if releasing failed and the result doesn't have an error yet, set
# an error
if (
exc_str is not None
and result is not None
and result.status != NodeStatus.Error
and error is None
):
error = exc_str
if error is not None:
result = self.error_result(ctx.node, error, started, ctx.timing)
elif result is not None:
result = self.from_run_result(result, started, ctx.timing)
else:
result = self.ephemeral_result(ctx.node, started, ctx.timing)
return result
def _safe_release_connection(self):
"""Try to release a connection. If an exception is hit, log and return
the error string.
"""
try:
self.adapter.release_connection()
except Exception as exc:
fire_event(
NodeConnectionReleaseError(
node_name=self.node.name, exc=str(exc), exc_info=traceback.format_exc()
)
)
return str(exc)
return None
def before_execute(self):
raise NotImplementedError()
def execute(self, compiled_node, manifest):
raise NotImplementedError()
def run(self, compiled_node, manifest):
return self.execute(compiled_node, manifest)
def after_execute(self, result):
raise NotImplementedError()
def _skip_caused_by_ephemeral_failure(self):
if self.skip_cause is None or self.skip_cause.node is None:
return False
return self.skip_cause.node.is_ephemeral_model
def on_skip(self):
schema_name = self.node.schema
node_name = self.node.name
error_message = None
if not self.node.is_ephemeral_model:
# if this model was skipped due to an upstream ephemeral model
# failure, print a special 'error skip' message.
if self._skip_caused_by_ephemeral_failure():
fire_event(
LogSkipBecauseError(
schema=schema_name,
relation=node_name,
index=self.node_index,
total=self.num_nodes,
)
)
print_run_result_error(result=self.skip_cause, newline=False)
if self.skip_cause is None: # mypy appeasement
raise DbtInternalError(
"Skip cause not set but skip was somehow caused by an ephemeral failure"
)
# set an error so dbt will exit with an error code
error_message = (
"Compilation Error in {}, caused by compilation error "
"in referenced ephemeral model {}".format(
self.node.unique_id, self.skip_cause.node.unique_id
)
)
else:
# 'skipped' nodes should not have a value for 'node_finished_at'
# they do have 'node_started_at', which is set in GraphRunnableTask.call_runner
self.node.update_event_status(node_status=RunStatus.Skipped)
fire_event(
SkippingDetails(
resource_type=self.node.resource_type,
schema=schema_name,
node_name=node_name,
index=self.node_index,
total=self.num_nodes,
node_info=self.node.node_info,
)
)
node_result = self.skip_result(self.node, error_message)
return node_result
def do_skip(self, cause=None):
self.skip = True
self.skip_cause = cause