From 27caf413f2d601de8d81d6cc3fb2614ff1bf354c Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Wed, 5 Feb 2025 14:24:15 +0100
Subject: [PATCH] Add support for extras to the labs CLI (#2288)

## Changes

Added support for extras / optional Python dependencies in the labs CLI.

Added new `extras` field under install.

Example:

```yaml
install:
  script: install.py
  extras: cli
```

Resolves: #2257

## Tests

Manual test
---
 cmd/labs/project/installer.go                                | 5 +++++
 cmd/labs/project/schema.json                                 | 5 +++++
 .../.databricks/labs/blueprint/lib/labs.yml                  | 1 +
 3 files changed, 11 insertions(+)

diff --git a/cmd/labs/project/installer.go b/cmd/labs/project/installer.go
index 05f7d68aa..2e42ce43d 100644
--- a/cmd/labs/project/installer.go
+++ b/cmd/labs/project/installer.go
@@ -32,6 +32,7 @@ type hook struct {
 	RequireDatabricksConnect bool    `yaml:"require_databricks_connect,omitempty"`
 	MinRuntimeVersion        string  `yaml:"min_runtime_version,omitempty"`
 	WarehouseTypes           whTypes `yaml:"warehouse_types,omitempty"`
+	Extras                   string  `yaml:"extras,omitempty"`
 }
 
 func (h *hook) RequireRunningCluster() bool {
@@ -258,6 +259,10 @@ func (i *installer) setupPythonVirtualEnvironment(ctx context.Context, w *databr
 		}
 	}
 	feedback <- "Installing Python library dependencies"
+	if i.Installer.Extras != "" {
+		// install main and optional dependencies
+		return i.installPythonDependencies(ctx, fmt.Sprintf(".[%s]", i.Installer.Extras))
+	}
 	return i.installPythonDependencies(ctx, ".")
 }
 
diff --git a/cmd/labs/project/schema.json b/cmd/labs/project/schema.json
index a779b15e4..7aa65813c 100644
--- a/cmd/labs/project/schema.json
+++ b/cmd/labs/project/schema.json
@@ -42,6 +42,11 @@
                 },
                 "warehouse_types": {
                     "enum": [ "PRO", "CLASSIC", "TYPE_UNSPECIFIED" ]
+                },
+                "extras": {
+                    "type": "string",
+                    "pattern": "^([^,]+)(,([^,]+))*$",
+                    "default": ""
                 }
             }
         },
diff --git a/cmd/labs/project/testdata/installed-in-home/.databricks/labs/blueprint/lib/labs.yml b/cmd/labs/project/testdata/installed-in-home/.databricks/labs/blueprint/lib/labs.yml
index 0ac4bf826..b8a0e695e 100644
--- a/cmd/labs/project/testdata/installed-in-home/.databricks/labs/blueprint/lib/labs.yml
+++ b/cmd/labs/project/testdata/installed-in-home/.databricks/labs/blueprint/lib/labs.yml
@@ -8,6 +8,7 @@ install:
   warehouse_types:
     - PRO
   script: install.py
+  extras: ""
 entrypoint: main.py
 min_python: 3.9
 commands: