Add multi schema support for databricks

README.md

@@ -630,7 +630,9 @@ Required parameters:
 - `your_databricks_workspace_id`: databricks workspace ID
 - `your_warehouse_id`: databricks sql warehouse ID
 - `your_catalog`: Unity Catalog name
-- `your_schema`: Schema/database name within the catalog
+
+optional parameters:
+- `your_schema`: Schema/database name within the catalog. If not specified all tables in the catalogue will be documented and generated table references will include schema name.
 
 Authentication (choose one):
 - PAT Authentication: `token` - Personal access token for authentication

datasource/databricks.go

@@ -24,9 +24,6 @@ func AnalyzeDatabricks(urlstr string) (_ *schema.Schema, err error) {
 		return nil, errors.New("no catalog name in the connection string")
 	}
 	schemaName := u.Query().Get("schema")
-	if schemaName == "" {
-		return nil, errors.New("no schema name in the connection string")
-	}
 
 	// Extract authentication parameters
 	token := u.Query().Get("token")
@@ -53,19 +50,15 @@ func AnalyzeDatabricks(urlstr string) (_ *schema.Schema, err error) {
 		return nil, errors.New("incomplete OAuth credentials: 'client_id' is required when 'client_secret' is provided")
 	}
 
-	s.Name = fmt.Sprintf("%s.%s", catalog, schemaName)
-
-	// Build databricks driver DSN based on authentication method
-	var databricksDSN string
-	if hasToken {
-		// PAT token authentication: token:TOKEN@host:port/path?catalog=CATALOG&schema=SCHEMA
-		databricksDSN = fmt.Sprintf("token:%s@%s%s?catalog=%s&schema=%s", token, u.Host, u.Path, catalog, schemaName)
+	if schemaName != "" {
+		s.Name = fmt.Sprintf("%s.%s", catalog, schemaName)
 	} else {
-		// OAuth client credentials authentication: host:port/path?authType=OauthM2M&clientID=ID&clientSecret=SECRET&catalog=CATALOG&schema=SCHEMA
-		databricksDSN = fmt.Sprintf("%s%s?authType=OauthM2M&clientID=%s&clientSecret=%s&catalog=%s&schema=%s",
-			u.Host, u.Path, clientID, clientSecret, catalog, schemaName)
+		s.Name = catalog
 	}
 
+	// Build databricks driver DSN based on authentication method
+	databricksDSN := buildDatabricksDSN(u.Host, u.Path, catalog, schemaName, token, clientID, clientSecret)
+
 	db, err := sql.Open("databricks", databricksDSN)
 	if err != nil {
 		return nil, err
@@ -75,9 +68,25 @@ func AnalyzeDatabricks(urlstr string) (_ *schema.Schema, err error) {
 	}()
 
 	driver := databricks.New(db)
+	if schemaName != "" {
+		driver.SetExplicitSchema(true)
+	}
 	if err := driver.Analyze(s); err != nil {
 		return nil, err
 	}
 
 	return s, nil
 }
+
+func buildDatabricksDSN(host, path, catalog, schema, token, clientID, clientSecret string) string {
+	baseParams := fmt.Sprintf("catalog=%s", catalog)
+	if schema != "" {
+		baseParams = fmt.Sprintf("%s&schema=%s", baseParams, schema)
+	}
+
+	if token != "" {
+		return fmt.Sprintf("token:%s@%s%s?%s", token, host, path, baseParams)
+	}
+	return fmt.Sprintf("%s%s?authType=OauthM2M&clientID=%s&clientSecret=%s&%s",
+		host, path, clientID, clientSecret, baseParams)
+}

datasource/databricks_test.go

@@ -0,0 +1,88 @@
+package datasource
+
+import (
+	"testing"
+)
+
+func TestBuildDatabricksDSN(t *testing.T) {
+	tests := []struct {
+		name         string
+		host         string
+		path         string
+		catalog      string
+		schema       string
+		token        string
+		clientID     string
+		clientSecret string
+		want         string
+	}{
+		{
+			name:    "token auth with schema",
+			host:    "host:443",
+			path:    "/sql/1.0/warehouses/abc123",
+			catalog: "my_catalog",
+			schema:  "my_schema",
+			token:   "dapi1234567890",
+			want:    "token:dapi1234567890@host:443/sql/1.0/warehouses/abc123?catalog=my_catalog&schema=my_schema",
+		},
+		{
+			name:    "token auth without schema",
+			host:    "host:443",
+			path:    "/sql/1.0/warehouses/abc123",
+			catalog: "my_catalog",
+			schema:  "",
+			token:   "dapi1234567890",
+			want:    "token:dapi1234567890@host:443/sql/1.0/warehouses/abc123?catalog=my_catalog",
+		},
+		{
+			name:         "oauth auth with schema",
+			host:         "host:443",
+			path:         "/sql/1.0/warehouses/abc123",
+			catalog:      "my_catalog",
+			schema:       "my_schema",
+			token:        "",
+			clientID:     "client123",
+			clientSecret: "secret456",
+			want:         "host:443/sql/1.0/warehouses/abc123?authType=OauthM2M&clientID=client123&clientSecret=secret456&catalog=my_catalog&schema=my_schema",
+		},
+		{
+			name:         "oauth auth without schema",
+			host:         "host:443",
+			path:         "/sql/1.0/warehouses/abc123",
+			catalog:      "my_catalog",
+			schema:       "",
+			token:        "",
+			clientID:     "client123",
+			clientSecret: "secret456",
+			want:         "host:443/sql/1.0/warehouses/abc123?authType=OauthM2M&clientID=client123&clientSecret=secret456&catalog=my_catalog",
+		},
+		{
+			name:    "token auth with different catalog",
+			host:    "dbc-123.cloud.databricks.com:443",
+			path:    "/sql/1.0/warehouses/xyz",
+			catalog: "production",
+			schema:  "analytics",
+			token:   "token_abc",
+			want:    "token:[email protected]:443/sql/1.0/warehouses/xyz?catalog=production&schema=analytics",
+		},
+		{
+			name:         "oauth auth with different catalog",
+			host:         "dbc-123.cloud.databricks.com:443",
+			path:         "/sql/1.0/warehouses/xyz",
+			catalog:      "production",
+			schema:       "analytics",
+			token:        "",
+			clientID:     "oauth_client",
+			clientSecret: "oauth_secret",
+			want:         "dbc-123.cloud.databricks.com:443/sql/1.0/warehouses/xyz?authType=OauthM2M&clientID=oauth_client&clientSecret=oauth_secret&catalog=production&schema=analytics",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := buildDatabricksDSN(tt.host, tt.path, tt.catalog, tt.schema, tt.token, tt.clientID, tt.clientSecret)
+			if got != tt.want {
+				t.Errorf("buildDatabricksDSN() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}