如何实现同时兼容Databricks与databricks-connect的Notebook导入？

这确实是个挺头疼的痛点——想一套代码通吃VS Code本地开发（用databricks-connect）和Databricks集群运行，不用来回改导入逻辑。我整理了几个实用的方案，你可以根据自己的场景选：

方案一：环境判断动态切换导入方式 #

核心思路是通过Spark配置判断当前运行环境，自动选择用import（VS Code+databricks-connect）还是%run（Databricks集群）。

实现代码示例 #

import sys
from pyspark.sql import SparkSession

def import_notebook(notebook_path):
    spark = SparkSession.getActiveSession()
    # 判断是否处于databricks-connect环境
    is_connect = spark.conf.get("spark.databricks.service.client.enabled", "false") == "true"
    
    if is_connect:
        # 处理VS Code环境：把Notebook路径转成Python模块名，添加到sys.path
        # 假设你的Notebook都放在/Workspace/Notebooks目录下
        notebook_dir = "/".join(notebook_path.split("/")[:-1])
        sys.path.append(notebook_dir)
        # 把Notebook文件名的横杠换成下划线（Python模块名不能有横杠）
        module_name = notebook_path.split("/")[-1].replace("-", "_")
        # 导入模块并执行需要的逻辑（比如调用里面的初始化函数）
        module = __import__(module_name)
        # 如果你的config-notebook里有init_config()这类函数，直接调用
        if hasattr(module, "init_config"):
            module.init_config()
    else:
        # Databricks集群环境：用exec执行%run魔法命令
        exec(f'%run "{notebook_path}"')

# 使用示例：导入config-notebook
import_notebook("/Workspace/Notebooks/config-notebook")

优缺点 #

• ✅ 不用大规模重构现有Notebook
• ❌ 需要处理路径转模块名的逻辑，还要确保Notebook所在目录在Python路径中
• ❌ 如果Notebook里是零散的代码（不是封装成函数），可能需要调整代码结构才能在import后正常执行

方案二：将可复用代码抽成Python脚本（最推荐） #

Notebook本质是交互式分析工具，可复用的配置、工具函数其实更适合放在.py脚本里。这样不管是VS Code还是Databricks，都能直接用import导入，完全没有兼容性问题。

操作步骤 #

1. 把config-notebook里的可复用代码（比如配置变量、工具函数）抽出来，保存为config.py（放在Workspace的同一个目录下）
2. 在需要用到的Notebook里，直接用import config导入
3. 不管是在VS Code用databricks-connect运行，还是在Databricks Web IDE里运行，这个导入语句都能正常工作

优缺点 #

• ✅ 完全符合Python模块化规范，代码结构更清晰
• ✅ 零兼容性问题，不用写任何环境判断逻辑
• ❌ 需要少量重构现有代码，把零散的Notebook代码封装成函数/类

方案三：用Databricks API调用Notebook（适合批量运行场景） #

如果你的需求是运行整个Notebook并获取结果，而不是导入里面的变量/函数，可以用Databricks Workflow API来实现。databricks-connect环境下可以通过databricks-sdk调用API运行Notebook。

代码示例 #

from databricks.sdk import WorkspaceClient
from databricks.sdk.service import jobs

def run_notebook_via_api(notebook_path):
    w = WorkspaceClient()
    # 创建临时任务运行Notebook
    run = w.jobs.create_run(
        tasks=[
            jobs.Task(
                task_key="run_config_notebook",
                notebook_task=jobs.NotebookTask(notebook_path=notebook_path)
            )
        ]
    )
    # 等待任务完成（可选）
    w.jobs.wait_run(run_id=run.run_id)
    # 获取运行结果（根据需要处理）
    run_details = w.jobs.get_run(run_id=run.run_id)
    return run_details

优缺点 #

• ✅ 适合批量调度或运行整个Notebook的场景
• ❌ 无法直接导入Notebook里的变量/函数，只能获取运行结果
• ❌ 配置稍复杂，需要确保API权限足够

内容的提问来源于stack exchange，提问作者Maxime