ssh: register the Databricks %pip magic on serverless environment 5+

.nextchanges/cli/ssh-jupyter-pip-magics.md

@@ -0,0 +1 @@
+* Fix `%pip` in Jupyter notebooks opened over `databricks ssh connect` on serverless environment versions 5 and 6 and on dedicated clusters running recent Databricks Runtime versions, which fell back to the stock IPython magic instead of the Databricks one. ([#7263](https://github.com/databricks/cli/pull/7263))

experimental/ssh/internal/server/jupyter-init.py

@@ -151,25 +151,53 @@ def _register_common_magics():
 @_log_exceptions
 def _register_pip_magics():
     """Register the pip magic command parser with IPython."""
+    import inspect
+
     from dbruntime import UserNamespaceInitializer
     from dbruntime.DatasetInfo import UserNamespaceDict
     from dbruntime.PipMagicOverrides import PipMagicOverrides
 
     user_namespace_initializer = UserNamespaceInitializer.getOrCreate()
     entry_point = user_namespace_initializer.get_spark_entry_point()
-    user_ns = UserNamespaceDict(
-        user_namespace_initializer.get_namespace_globals(),
-        entry_point.getDriverConf(),
-        entry_point,
-    )
     ip = get_ipython()
 
-    try:
-        # Older DBRs
-        pip_magic = PipMagicOverrides(entry_point, ip.user_ns["sc"]._conf, user_ns)
-    except Exception:
-        # Newer DBRs
-        pip_magic = PipMagicOverrides(entry_point, user_ns, ip)
+    params = inspect.signature(PipMagicOverrides).parameters
+    if "logger" in params:
+        # Runtimes that take a logger (serverless environment 5+ and recent Databricks Runtime
+        # versions) are migrating off Py4J, so build the magic from the driver connection the way
+        # dbruntime's own kernel does: on serverless environment 6 the entry point is None. The
+        # constructor gained keyword arguments over several runtime releases, so pass only the ones
+        # this runtime accepts. Values are looked up lazily because the attributes behind
+        # newer arguments may not exist on runtimes that don't accept them.
+        driver_connection = user_namespace_initializer.get_driver_connection()
+        driver_conf = driver_connection.runtime_integration_client.getDriverConf()
+        user_ns = UserNamespaceDict(user_namespace_initializer.get_namespace_globals(), driver_conf, entry_point)
+        optional_kwargs = {
+            "orig_pip_magic": lambda: ip.magics_manager.magics["line"]["pip"],
+            "orig_conda_magic": lambda: ip.magics_manager.magics["line"]["conda"],
+            "driver_conf": lambda: driver_conf,
+            "environment_resource_client": lambda: driver_connection.environment_resource_client,
+            "sc_conf": lambda: user_namespace_initializer.db_connection.sc_conf,
+        }
+        pip_magic = PipMagicOverrides(
+            entry_point,
+            user_ns,
+            ip,
+            logger=driver_connection.observability_client.getLogger(),
+            **{name: get_value() for name, get_value in optional_kwargs.items() if name in params},
+        )
+    else:
+        user_ns = UserNamespaceDict(
+            user_namespace_initializer.get_namespace_globals(),
+            entry_point.getDriverConf(),
+            entry_point,
+        )
+        try:
+            # Older DBRs
+            pip_magic = PipMagicOverrides(entry_point, ip.user_ns["sc"]._conf, user_ns)
+        except Exception:
+            # Newer DBRs
+            pip_magic = PipMagicOverrides(entry_point, user_ns, ip)
 
     ip.register_magics(pip_magic)