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 | """WsfsImportHook Patch for Editable Installs.
This module monkey-patches workspace import machinery to allow imports from
editable install paths. Supports both legacy and modern Databricks runtimes:
- DBR < 18.0: Patches dbruntime.wsfs_import_hook.WsfsImportHook
- DBR >= 18.0: Patches dbruntime.workspace_import_machinery._WorkspacePathEntryFinder
The import hooks normally block imports that don't originate from:
- site-packages
- /Workspace paths
- Whitelisted paths
This patch extends the whitelist to include editable install paths detected
from .pth files and .egg-link files.
"""
from collections.abc import Callable
import inspect
import sys
from typing import Any
from dbx_patch.base_patch import BasePatch
from dbx_patch.models import PatchResult
from dbx_patch.utils.runtime_version import is_runtime_version_gte
class WsfsImportHookPatch(BasePatch):
"""Patch for workspace import machinery.
Patches the appropriate import hook based on runtime version to allow
imports from editable install paths.
"""
def _create_patched_is_user_import_legacy(self, original_method: Callable[..., bool]) -> Callable[..., bool]:
"""Create patched version for WsfsImportHook (DBR < 18.0).
Args:
original_method: The original __is_user_import method
Returns:
Patched method that includes editable path checking
"""
def patched_is_user_import(hook_self: Any) -> bool:
logger = self._get_logger()
if logger:
logger.debug("WsfsImportHook.__is_user_import called (PATCHED)")
try:
f = inspect.currentframe()
num_items_processed = 0
while f is not None:
# Prevent infinite loops
if num_items_processed >= hook_self._WsfsImportHook__max_recursion_depth:
return True
filename = hook_self.get_filename(f)
# Allow whitelisted paths (existing behavior)
allow_import = any(
whitelisted_item in filename for whitelisted_item in hook_self.SITE_PACKAGE_WHITE_LIST
)
if allow_import:
return True
# NEW: Allow imports from editable install paths
if self._cached_editable_paths:
is_editable_path = any(
filename.startswith(editable_path) for editable_path in self._cached_editable_paths
)
if is_editable_path and logger:
matching = [p for p in self._cached_editable_paths if filename.startswith(p)]
logger.debug(
f"WsfsImportHook: Allowing import from editable path: {filename} (matched: {matching[0] if matching else 'unknown'})"
)
return True
# Check if from site-packages (existing behavior)
is_site_packages = any(
filename.startswith(package) for package in hook_self._WsfsImportHook__site_packages
)
if is_site_packages:
return False
num_items_processed += 1
f = f.f_back
# None of the stack frames are from site-packages, probably from user
return True
except Exception as e:
# Fail open - allow the import if we can't determine
logger = self._get_logger()
if logger:
logger.warning(f"DBX-Patch: Exception in patched __is_user_import: {e}")
return False
return patched_is_user_import
def _create_patched_is_user_import_modern(self, original_method: Callable[..., bool]) -> Callable[..., bool]:
"""Create patched version for _WorkspacePathEntryFinder (DBR >= 18.0).
Args:
original_method: The original _is_user_import method
Returns:
Patched method that includes editable path checking
"""
def patched_is_user_import(finder_self: Any) -> bool:
logger = self._get_logger()
if logger:
logger.debug("_WorkspacePathEntryFinder._is_user_import called (PATCHED)")
try:
frame = sys._getframe()
num_items_processed = 0
while frame is not None:
if num_items_processed >= finder_self._max_stack_depth:
return True
filename = finder_self.get_filename(frame)
# NEW: Allow imports from editable install paths FIRST
if self._cached_editable_paths:
is_editable_path = any(
filename.startswith(editable_path) for editable_path in self._cached_editable_paths
)
if is_editable_path and logger:
matching = [p for p in self._cached_editable_paths if filename.startswith(p)]
logger.debug(
f"_WorkspacePathEntryFinder: Allowing import from editable path: {filename} (matched: {matching[0] if matching else 'unknown'})"
)
return True
# Check allow list (existing behavior)
if any(allow_listed_item in filename for allow_listed_item in finder_self.SITE_PACKAGE_ALLOW_LIST):
return True
# Check if from site-packages (existing behavior)
if any(filename.startswith(package) for package in finder_self._site_packages):
return False
num_items_processed += 1
frame = frame.f_back
return True
except Exception as e:
logger = self._get_logger()
if logger:
logger.warning(f"DBX-Patch: Exception in patched _is_user_import: {e}")
return True
return patched_is_user_import
def patch(self) -> PatchResult:
"""Apply the workspace import hook patch.
Automatically detects runtime version and applies appropriate patch.
Returns:
PatchResult with operation details
"""
logger = self._get_logger()
if self._is_applied:
if logger:
logger.info("Workspace import hook patch already applied.")
return PatchResult(
success=True,
already_patched=True,
editable_paths_count=len(self._cached_editable_paths),
editable_paths=sorted(self._cached_editable_paths),
hook_found=True,
)
try:
# Detect editable paths
self._cached_editable_paths = self._detect_editable_paths()
# Determine which version to patch based on runtime version
use_modern = is_runtime_version_gte(18, 0)
if use_modern:
# Patch modern _WorkspacePathEntryFinder (DBR >= 18.0)
if logger:
logger.info("Detected DBR >= 18.0, using modern workspace_import_machinery patch")
try:
from dbruntime.workspace_import_machinery import ( # type: ignore[import-not-found]
_WorkspacePathEntryFinder,
)
if logger:
logger.info(
f"Patching modern _WorkspacePathEntryFinder to allow {len(self._cached_editable_paths)} editable install path(s)..."
)
# Save original method
self._original_target = _WorkspacePathEntryFinder._is_user_import
if self._original_target is None:
if logger:
logger.error("Failed to save original method")
return PatchResult(
success=False,
already_patched=False,
editable_paths_count=0,
editable_paths=[],
hook_found=True,
error="Failed to save original method",
)
# Create and apply patch
patched_method = self._create_patched_is_user_import_modern(self._original_target)
_WorkspacePathEntryFinder._is_user_import = patched_method
if logger:
logger.success("Modern _WorkspacePathEntryFinder patched successfully!")
result = PatchResult(
success=True,
already_patched=False,
editable_paths_count=len(self._cached_editable_paths),
editable_paths=sorted(self._cached_editable_paths),
hook_found=True,
)
except ImportError as e:
result = PatchResult(
success=False,
already_patched=False,
editable_paths_count=0,
editable_paths=[],
hook_found=False,
error=f"Modern module not found: {e}",
)
else:
# Patch legacy WsfsImportHook (DBR < 18.0)
if logger:
logger.info("Detected DBR < 18.0, using legacy WsfsImportHook patch")
try:
from dbruntime.wsfs_import_hook import WsfsImportHook # type: ignore[import-not-found]
if logger:
logger.info(
f"Patching legacy WsfsImportHook to allow {len(self._cached_editable_paths)} editable install path(s)..."
)
# Save original method
self._original_target = WsfsImportHook._WsfsImportHook__is_user_import
if self._original_target is None:
if logger:
logger.error("Failed to save original method")
return PatchResult(
success=False,
already_patched=False,
editable_paths_count=0,
editable_paths=[],
hook_found=True,
error="Failed to save original method",
)
# Create and apply patch
patched_method = self._create_patched_is_user_import_legacy(self._original_target)
WsfsImportHook._WsfsImportHook__is_user_import = patched_method
if logger:
logger.success("Legacy WsfsImportHook patched successfully!")
result = PatchResult(
success=True,
already_patched=False,
editable_paths_count=len(self._cached_editable_paths),
editable_paths=sorted(self._cached_editable_paths),
hook_found=True,
)
except ImportError as e:
result = PatchResult(
success=False,
already_patched=False,
editable_paths_count=0,
editable_paths=[],
hook_found=False,
error=f"Legacy module not found: {e}",
)
if result.success:
self._is_applied = True
if self._cached_editable_paths and logger:
with logger.indent():
logger.info("Allowed editable paths:")
for path in sorted(self._cached_editable_paths):
logger.info(f" - {path}")
elif logger:
with logger.indent():
logger.warning("No editable install paths found yet.")
logger.info("Run 'pip install -e .' first, then reapply patches.")
return result
except ImportError as e:
if logger:
logger.warning(f"Could not import workspace import machinery: {e}")
with logger.indent():
logger.info("This is normal if not running in Databricks environment.")
return PatchResult(
success=False,
already_patched=False,
editable_paths_count=0,
editable_paths=[],
hook_found=False,
error=str(e),
)
except Exception as e:
if logger:
logger.error(f"Error patching workspace import machinery: {e}")
return PatchResult(
success=False,
already_patched=False,
editable_paths_count=0,
editable_paths=[],
hook_found=True,
error=str(e),
)
def remove(self) -> bool:
"""Remove the patch and restore original workspace import hook behavior.
Automatically detects which runtime version was patched and restores accordingly.
Returns:
True if unpatch was successful, False otherwise
"""
logger = self._get_logger()
if not self._is_applied:
if logger:
logger.info("No patch to remove.")
return False
if self._original_target is None:
if logger:
logger.warning("Original method not saved, cannot unpatch.")
return False
try:
# Determine which version was patched based on runtime version
use_modern = is_runtime_version_gte(18, 0)
if use_modern:
# Restore modern _WorkspacePathEntryFinder (DBR >= 18.0)
try:
from dbruntime.workspace_import_machinery import ( # type: ignore[import-not-found]
_WorkspacePathEntryFinder,
)
_WorkspacePathEntryFinder._is_user_import = self._original_target
self._is_applied = False
self._original_target = None
if logger:
logger.success("_WorkspacePathEntryFinder patch removed successfully.")
return True
except ImportError as e:
if logger:
logger.error(f"Error restoring modern patch: {e}") # noqa: TRY400
return False
else:
# Restore legacy WsfsImportHook (DBR < 18.0)
try:
from dbruntime.wsfs_import_hook import WsfsImportHook # type: ignore[import-not-found]
WsfsImportHook._WsfsImportHook__is_user_import = self._original_target
self._is_applied = False
self._original_target = None
if logger:
logger.success("WsfsImportHook patch removed successfully.")
return True
except ImportError as e:
if logger:
logger.error(f"Error restoring legacy patch: {e}") # noqa: TRY400
return False
except Exception as e:
if logger:
logger.error(f"Error removing patch: {e}") # noqa: TRY400
return False
def is_applied(self) -> bool:
"""Check if the workspace import hook patch is currently applied.
Returns:
True if patched, False otherwise
"""
return self._is_applied
|