Skip to content

Commit 89a5a9f

Browse files
authored
Fix databricks connector postprocessing (#85)
* Applied defensive programming to ensure connector registration is firing * Fixed broken process and kwargs stripping * flake8 fix * Added timeout * undo small change * Removed kwarg cleaner
1 parent b838205 commit 89a5a9f

File tree

1 file changed

+14
-7
lines changed

1 file changed

+14
-7
lines changed

noteable_magics/datasource_postprocessing.py

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
import shutil
33
from base64 import b64decode
44
from pathlib import Path
5+
from subprocess import PIPE, Popen
56
from tempfile import NamedTemporaryFile
67
from typing import Any, Callable, Dict
78
from urllib.parse import quote_plus, urlparse
@@ -295,13 +296,19 @@ def postprocess_databricks(
295296
if connect_file_path.exists():
296297
connect_file_path.unlink()
297298

298-
# Now let databricks-connect external command (re)build it and do whatever
299-
# else it does. See ENG-5517. The 'y' at the start accepts the license agreement.
300-
# (We've fallen oh so far from Don Libes' tcl Expect for stuff like this.)
301-
pipeline = f"echo y {args['host']} {args['token']} {args[cluster_id_key]} {args['org_id']} {args['port']} | databricks-connect configure"
302-
os.system(pipeline)
299+
p = Popen(['databricks-connect', 'configure'], stdout=PIPE, stdin=PIPE, stderr=PIPE)
300+
_stdout, stderr = p.communicate(input=f"""y
301+
{args['host']}
302+
{args['token']}
303+
{args[cluster_id_key]}
304+
{args['org_id']}
305+
{args['port']}""".encode(), timeout=10)
303306

304-
# Always be sure to purge these only-for-databricks-connect file args from create_engine_kwargs,
307+
if p.returncode != 0:
308+
# Failed to exectute the script. Raise an exception.
309+
raise ValueError("Failed to execute databricks-connect configure script: " + stderr)
310+
311+
# Always be sure to purge these only-for-databricks-connect file args from connect_args,
305312
# even if not all were present.
306313
for key in connect_file_opt_keys:
307-
create_engine_kwargs.pop(key, '')
314+
connect_args.pop(key, '')

0 commit comments

Comments
 (0)