Skip to content

Commit a2d7f14

Browse files
Add cwd and project root to sys.path for dbconnect runner (#1221)
## Changes * Sys.path by default did not contain the cwd and project root. This PR adds them from the bootstrap code. ## Tests <!-- How is this tested? -->
1 parent d8604d8 commit a2d7f14

File tree

1 file changed

+21
-11
lines changed

1 file changed

+21
-11
lines changed

packages/databricks-vscode/resources/python/dbconnect-bootstrap.py

Lines changed: 21 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
11
import os
2+
import sys
3+
import logging
4+
from runpy import run_path
25

36
# Load environment variables from .databricks/.databricks.env
47
# We only look for the folder in the current working directory
5-
# since for all commands laucnhed from
8+
# since for all commands laucnhed from root workspace
69
def load_env_file_from_cwd(path: str):
710
if not os.path.isdir(path):
811
return
@@ -17,17 +20,29 @@ def load_env_file_from_cwd(path: str):
1720
os.environ[key] = value
1821
return
1922

20-
load_env_file_from_cwd(os.getcwd())
23+
24+
script = sys.argv[1]
25+
sys.argv = sys.argv[1:]
26+
logging.debug(f"Running ${script}")
27+
logging.debug(f"args: ${sys.argv[1:]}")
28+
29+
try:
30+
cur_dir = os.path.dirname(script)
31+
except Exception as e:
32+
logging.error(f"Failed to get current directory: {e}")
33+
cur_dir = os.getcwd()
34+
35+
root_dir = os.getcwd()
36+
load_env_file_from_cwd(root_dir)
2137

2238
log_level = os.environ.get("DATABRICKS_VSCODE_LOG_LEVEL")
2339
log_level = log_level if log_level is not None else "WARN"
2440

25-
import logging
2641
logging.basicConfig(level=log_level)
2742

2843
db_globals = {}
2944

30-
from databricks.sdk.runtime import dbutils
45+
from databricks.sdk.runtime import dbutils # noqa: E402
3146
db_globals['dbutils'] = dbutils
3247

3348
# "table", "sc", "sqlContext" are missing
@@ -50,12 +65,7 @@ def getArgument(*args, **kwargs):
5065

5166
db_globals['getArgument'] = getArgument
5267

53-
from runpy import run_path
54-
import sys
55-
56-
script = sys.argv[1]
57-
sys.argv = sys.argv[1:]
58-
logging.debug(f"Running ${script}")
59-
logging.debug(f"args: ${sys.argv[1:]}")
68+
sys.path.insert(0, root_dir)
69+
sys.path.insert(0, cur_dir)
6070

6171
run_path(script, init_globals=db_globals, run_name="__main__")

0 commit comments

Comments
 (0)