Initial version of DBFS Scanner to help with DBFS disablement (#379)

For upcoming functionality of full DBFS disablement we need to allow
customers to analyze their workspace and identify DBFS use to store data
and in the compute resources.

See README for details

Author: alexott

dbfs-scanner/README.md

@@ -0,0 +1,38 @@
+# A workspace scanner in preparation to DBFS disablement
+
+This folder contains code that could help with preparing to complete disablement of DBFS.  There are two tools available (each tool has a command-line variant and as notebook with `_nb` suffix):
+
+* `scan_compute.py` (and `scan_compute_nb`) - performs scanning of compute resources (interactive clusters, jobs, DLT pipelines, policies) for use of resources on DBFS (init scripts, libraries, etc.) and output findings.
+* `scan_dbfs.py` (and `scan_dbfs_nb`) - performs scanning of DBFS content and trying to classify it and estimate size (Delta tables, Structure Streaming checkpoints, DLT pipelines, ...).  By default it scans the whole DBFS Root (user-accessible part), but you can pass an optional start directory as parameter.
+
+
+## Usage as command-line utility
+
+1. Install `databricks-sdk` Python package.
+1. Setup [environment variables with authentication parameters](https://docs.databricks.com/aws/en/dev-tools/auth/) for a workspace that will be analyzed.
+1. Run a specific tool:
+
+  1. `python scan_compute.py` will scan all compute resources and output results to console and also write them into `compute_scan_results.json` file.
+  1. `python scan_dbfs.py [start_directory]` will scan DBFS Root (or only a `start_directory` if specified), and output results to console and also write them into `dbfs_scan_results.json` file. 
+
+## Usage inside the workspace
+
+1. As **workspace administrator** open the corresponding notebook (`scan_compute_nb` or `scan_dbfs_nb`).
+1. Attach to a cluster (i.e. Serverless).
+1. Specify parameters in widgets (start directory for `scan_dbfs_nb`, or output directory if you want to persist result to UC Volume (`/Volumes/<catalog>/<schema>/<volume>/`).
+1. Press "Run all" and wait for finishing the execution.
+1. If output directory is specified, files `compute_scan_results.json` or `dbfs_scan_results.json` will be stored in that directory.
+
+
+## Known issues
+
+* Scan of DBFS is very slow due to the single-threaded implementation.
+* Output file names for results are hardcoded.
+
+
+## TODOs
+
+* \[ \] Use `blueprints` library for logging
+* \[ \] For all compute objects try to find last run/update timestamp...
+* \[ \] Parallelize scan of DBFS
+* \[ \] Allow to customize output file name in command line utilities

dbfs-scanner/helpers/__init__.py

@@ -0,0 +1,289 @@
+import base64
+import json
+import re
+from typing import Optional, Tuple
+
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.service.compute import ClusterSource, ClusterDetails, ClusterSpec
+from databricks.sdk.service.jobs import Task
+from databricks.sdk.service.pipelines import PipelineCluster
+from databricks.sdk.service.workspace import ExportFormat
+
+
+__notebook_pip_install_from_dbfs__ = re.compile(r"^#\s+MAGIC\s+%pip\s+install.*(/dbfs/.*)$")
+__notebook_dbfs_fuse_use__ = re.compile(r"^.*[\"'](/dbfs/[^\"']*)[\"'].*$")
+__notebook_dbfs_hdfs_use__ = re.compile(r"^.*[\"'](dbfs:/[^\"']*)[\"'].*$")
+
+
+def _analyze_notebook_or_wsfile(wc: WorkspaceClient, notebook_path: str) -> Tuple[list, list]:
+    libs = []
+    dbfs_file_refs = []
+    try:
+        b64 = wc.workspace.export(notebook_path, format=ExportFormat.SOURCE).content
+        content = base64.b64decode(b64).decode("utf-8")
+        for l in content.split("\n"):
+            m = __notebook_pip_install_from_dbfs__.match(l)
+            if m:
+                libs.append(m.group(1))
+            elif not l.lstrip().startswith("#"):
+                m = __notebook_dbfs_fuse_use__.match(l)
+                if m:
+                    dbfs_file_refs.append(m.group(1))
+                m = __notebook_dbfs_hdfs_use__.match(l)
+                if m:
+                    dbfs_file_refs.append(m.group(1))
+
+    except Exception as e:
+        # print(f"Error occurred while analyzing notebook {notebook_path}: {e}")
+        pass
+
+    return libs, dbfs_file_refs
+
+
+def analyze_dlt_pipelines(wc: WorkspaceClient) -> dict:
+    finds = {}
+    i = 0
+    for l in wc.pipelines.list_pipelines(max_results=100):
+        i += 1
+        if i % 100 == 0:
+            print(f"Scanned {i} DLT pipelines")
+        # print("Analyzing pipeline:", l.pipeline_id, l.name)
+        p = wc.pipelines.get(l.pipeline_id)
+        if not p:
+            continue
+        p_finds = {}
+        # analyze clusters
+        clusters = {}
+        for cl in p.spec.clusters:
+            cl_finds = _analyze_cluster_spec(cl, {})
+            if cl_finds:
+                clusters[cl.label] = cl_finds
+        if clusters:
+            p_finds["clusters"] = clusters
+        # analyze storage
+        if p.spec.storage:
+            p_finds["storage"] = p.spec.storage
+        # analyze libraries
+        lib_finds = {}
+        for lib in p.spec.libraries:
+            lib_path = ""
+            if lib.notebook:
+                lib_path = lib.notebook.path
+            elif lib.file:
+                lib_path = lib.file.path
+            if lib_path:
+                libs, dbfs_file_refs = _analyze_notebook_or_wsfile(wc, lib_path)
+                if libs or dbfs_file_refs:
+                    d = {}
+                    if libs:
+                        d["libraries"] = libs
+                    if dbfs_file_refs:
+                        d["dbfs_file_refs"] = dbfs_file_refs
+                    lib_finds[lib_path] = d
+        if lib_finds:
+            p_finds["libraries"] = lib_finds
+        if p_finds:
+            p_finds["name"] = l.name
+            finds[l.pipeline_id] = p_finds
+
+    print(f"Total {i} DLT pipelines")
+
+    return finds
+
+
+def _analyze_task(wc: WorkspaceClient, task: Task ) -> dict:
+    finds = {}
+    if task.spark_python_task:
+        if task.spark_python_task.python_file.startswith("dbfs:/"):
+            finds["python_file_on_dbfs"] = task.spark_python_task.python_file
+        elif task.spark_python_task.python_file.startswith("/"):
+            libs, dbfs_file_refs = _analyze_notebook_or_wsfile(wc, task.spark_python_task.python_file)
+            if libs or dbfs_file_refs:
+                finds["python_file_in_ws"] = {
+                    "path": task.spark_python_task.python_file
+                }
+                if libs:
+                    finds["python_file_in_ws"]["libraries"] = libs
+                if dbfs_file_refs:
+                    finds["python_file_in_ws"]["dbfs_file_refs"] = dbfs_file_refs
+
+    if task.notebook_task:
+        libs, dbfs_file_refs = _analyze_notebook_or_wsfile(wc, task.notebook_task.notebook_path)
+        if libs or dbfs_file_refs:
+            finds["notebook"] = {
+                "path": task.notebook_task.notebook_path,
+            }
+            if libs:
+                finds["notebook"]["libraries"] = libs
+            if dbfs_file_refs:
+                finds["notebook"]["dbfs_file_refs"] = dbfs_file_refs
+
+    if task.for_each_task:
+        fe_finds = _analyze_task(wc, task.for_each_task.task)
+        if fe_finds:
+            finds["for_each_task"] = fe_finds
+
+    for lib in (task.libraries or []):
+        dbfs_lib = ""
+        if lib.jar and lib.jar.startswith("dbfs:/"):
+            dbfs_lib = lib.jar
+        if lib.whl and lib.whl.startswith("dbfs:/"):
+            dbfs_lib = lib.whl
+        if lib.egg and lib.egg.startswith("dbfs:/"):
+            dbfs_lib = lib.egg
+        if dbfs_lib:
+            r = finds.get("libraries", [])
+            r.append(dbfs_lib)
+            finds["libraries"] = r
+
+    if task.new_cluster:
+        finds = _analyze_cluster_spec(task.new_cluster, finds)
+
+    return finds
+
+
+def analyze_jobs(wc: WorkspaceClient) -> dict:
+    res = {}
+    i = 0
+    for job in wc.jobs.list(expand_tasks=True, limit=100):
+        i += 1
+        if i % 100 == 0:
+            print(f"Scanned {i} jobs")
+        finds = {}
+        js = job.settings
+        # print("Analyzing job:", js.name)
+        for task in (js.tasks or []):
+            task_res = _analyze_task(wc, task)
+
+            if task_res:
+                t = finds.get("tasks", {})
+                t[task.task_key] = task_res
+                finds["tasks"] = t
+
+        jcs_finds = {}
+        for jc in (js.job_clusters or []):
+            jc_finds = {}
+            if jc.new_cluster:
+                jc_finds = _analyze_cluster_spec(jc.new_cluster, jc_finds)
+            if jc_finds:
+                jcs_finds[jc.job_cluster_key] = jc_finds
+
+        if jcs_finds:
+            finds["job_clusters"] = jcs_finds
+
+        if finds:
+            finds["job_name"] = js.name
+            res[job.job_id] = finds
+
+    print(f"Total {i} jobs")
+    return res
+
+
+def _analyze_cluster_spec(cl: ClusterDetails | ClusterSpec | PipelineCluster, finds: dict):
+    # check if we have any init scripts pointing to DBFS
+    for init_script in (cl.init_scripts or []):
+        if init_script.dbfs:
+            r = finds.get("init_scripts", [])
+            r.append(init_script.dbfs.destination)
+            finds["init_scripts"] = r
+    # check if we have cluster conf pointing to DBFS
+    if cl.cluster_log_conf and cl.cluster_log_conf.dbfs and cl.cluster_log_conf.dbfs.destination:
+        finds["cluster_log_conf"] = cl.cluster_log_conf.dbfs.destination
+
+    return finds
+
+
+def analyze_clusters(wc: WorkspaceClient) -> dict:
+    res = {}
+    i = 0
+    for cl in wc.clusters.list(page_size=100):
+        i += 1
+        if i % 100 == 0:
+            print(f"Scanned {i} clusters")
+        if cl.cluster_source not in [ClusterSource.UI, ClusterSource.API]:
+            continue
+        # print("Analyzing cluster:", cl.cluster_name)
+        finds = {}
+
+        # check if we have any libraries pointing to DBFS
+        for lib in wc.libraries.cluster_status(cl.cluster_id):
+            dbfs_lib = ""
+            if lib.library.jar and lib.library.jar.startswith("dbfs:/"):
+                dbfs_lib = lib.library.jar
+            if lib.library.whl and lib.library.whl.startswith("dbfs:/"):
+                dbfs_lib = lib.library.whl
+            if dbfs_lib:
+                r = finds.get("libraries", [])
+                r.append(dbfs_lib)
+                finds["libraries"] = r
+
+        finds = _analyze_cluster_spec(cl, finds)
+
+        # if we found anything, add it to the results
+        if finds:
+            finds["cluster_name"] = cl.cluster_name
+            res[cl.cluster_id] = finds
+
+    print(f"Total {i} clusters")
+    return res
+
+
+def _check_policy_definition(sdef: Optional[str], finds: dict):
+    policy_def = json.loads(sdef or "{}")
+    for k, v in policy_def.items():
+        if not isinstance(v, dict):
+            continue
+        typ = v.get("type")
+        if not typ:
+            continue
+        val = v.get("value") or v.get("defaultValue")
+        if typ == "fixed" and k == "cluster_log_conf.path" and val and val.startswith("dbfs:/"):
+            finds["cluster_log_conf"] = val
+        if typ == "fixed" and k.startswith("init_scripts.") and k.endswith(".dbfs.destination"):
+            r = finds.get("init_scripts", [])
+            r.append(val)
+            finds["init_scripts"] = r
+
+    return finds
+
+
+def analyze_cluster_policies(wc: WorkspaceClient) -> dict:
+    res = {}
+    i = 0
+    for pf in wc.policy_families.list():
+        i += 1
+        if i % 100 == 0:
+            print(f"Scanned {i} policies")
+        finds = {}
+        # print("Analyzing cluster policy family:", pf.name)
+        finds = _check_policy_definition(pf.definition, finds)
+
+        if finds:
+            finds["policy_name"] = pf.name
+            res[pf.policy_id] = finds
+
+    for pl in wc.cluster_policies.list():
+        i += 1
+        if i % 100 == 0:
+            print(f"Scanned {i} policies")
+        # print("Analyzing cluster policy:", pl.name)
+        finds = {}
+        for lib in (pl.libraries or []):
+            dbfs_lib = ""
+            if lib.jar and lib.jar.startswith("dbfs:/"):
+                dbfs_lib = lib.jar
+            if lib.whl and lib.whl.startswith("dbfs:/"):
+                dbfs_lib = lib.whl
+            if dbfs_lib:
+                r = finds.get("libraries", [])
+                r.append(dbfs_lib)
+                finds["libraries"] = r
+
+        finds = _check_policy_definition(pl.definition, finds)
+
+        if finds:
+            finds["policy_name"] = pl.name
+            res[pl.policy_id] = finds
+
+    print(f"Total {i} cluster policies and families")
+    return res