Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update table names for databricks_workspace_workspace and databricks_catalog_catalog #4

Merged
merged 8 commits into from
Aug 28, 2023
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions databricks/plugin.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ func Plugin(ctx context.Context) *plugin.Plugin {
Schema: ConfigSchema,
},
TableMap: map[string]*plugin.Table{
"databricks_catalog_catalog": tableDatabricksCatalogCatalog(ctx),
"databricks_catalog": tableDatabricksCatalog(ctx),
"databricks_catalog_connection": tableDatabricksCatalogConnection(ctx),
"databricks_catalog_external_location": tableDatabricksCatalogExternalLocation(ctx),
"databricks_catalog_function": tableDatabricksCatalogFunction(ctx),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Even though it's making an exception to the rule of following the SDK directory naming for services (the most consistent source we could find), we could treat job and pipeline as the service names instead of jobs and pipelines respectively. This would mean we'd have the table names:

  • databricks_job
  • databricks_job_run
  • databricks_pipeline
  • databricks_pipeline_event
  • databricks_pipeline_update

These seem like better table names than what we originally have

Expand Down Expand Up @@ -73,7 +73,7 @@ func Plugin(ctx context.Context) *plugin.Plugin {
"databricks_workspace_repo": tableDatabricksWorkspaceRepo(ctx),
"databricks_workspace_scope": tableDatabricksWorkspaceScope(ctx),
"databricks_workspace_secret": tableDatabricksWorkspaceSecret(ctx),
"databricks_workspace_workspace": tableDatabricksWorkspaceWorkspace(ctx),
"databricks_workspace": tableDatabricksWorkspace(ctx),
},
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,17 +11,17 @@ import (

//// TABLE DEFINITION

func tableDatabricksCatalogCatalog(_ context.Context) *plugin.Table {
func tableDatabricksCatalog(_ context.Context) *plugin.Table {
return &plugin.Table{
Name: "databricks_catalog_catalog",
Name: "databricks_catalog",
Description: "Gets an array of catalogs in the metastore.",
List: &plugin.ListConfig{
Hydrate: listCatalogCatalogs,
Hydrate: listCatalogs,
},
Get: &plugin.GetConfig{
KeyColumns: plugin.SingleColumn("name"),
ShouldIgnoreError: isNotFoundError([]string{"CATALOG_DOES_NOT_EXIST"}),
Hydrate: getCatalogCatalog,
Hydrate: getCatalog,
},
Columns: databricksAccountColumns([]*plugin.Column{
{
Expand Down Expand Up @@ -141,7 +141,7 @@ func tableDatabricksCatalogCatalog(_ context.Context) *plugin.Table {
Name: "workspace_bindings",
Description: "Array of workspace bindings.",
Type: proto.ColumnType_JSON,
Hydrate: getCatalogCatalogWorkspaceBindings,
Hydrate: getCatalogWorkspaceBindings,
Transform: transform.FromValue(),
},

Expand All @@ -158,19 +158,19 @@ func tableDatabricksCatalogCatalog(_ context.Context) *plugin.Table {

//// LIST FUNCTION

func listCatalogCatalogs(ctx context.Context, d *plugin.QueryData, h *plugin.HydrateData) (interface{}, error) {
func listCatalogs(ctx context.Context, d *plugin.QueryData, h *plugin.HydrateData) (interface{}, error) {
logger := plugin.Logger(ctx)

// Create client
client, err := connectDatabricksWorkspace(ctx, d)
if err != nil {
logger.Error("databricks_catalog_catalog.listCatalogCatalogs", "connection_error", err)
logger.Error("databricks_catalog.listCatalogs", "connection_error", err)
return nil, err
}

catalogs, err := client.Catalogs.ListAll(ctx)
if err != nil {
logger.Error("databricks_catalog_catalog.listCatalogCatalogs", "api_error", err)
logger.Error("databricks_catalog.listCatalogs", "api_error", err)
return nil, err
}

Expand All @@ -188,7 +188,7 @@ func listCatalogCatalogs(ctx context.Context, d *plugin.QueryData, h *plugin.Hyd

//// HYDRATE FUNCTIONS

func getCatalogCatalog(ctx context.Context, d *plugin.QueryData, _ *plugin.HydrateData) (interface{}, error) {
func getCatalog(ctx context.Context, d *plugin.QueryData, _ *plugin.HydrateData) (interface{}, error) {
logger := plugin.Logger(ctx)
name := d.EqualsQualString("name")

Expand All @@ -200,33 +200,33 @@ func getCatalogCatalog(ctx context.Context, d *plugin.QueryData, _ *plugin.Hydra
// Create client
client, err := connectDatabricksWorkspace(ctx, d)
if err != nil {
logger.Error("databricks_catalog_catalog.getCatalogCatalog", "connection_error", err)
logger.Error("databricks_catalog.getCatalog", "connection_error", err)
return nil, err
}

catalog, err := client.Catalogs.GetByName(ctx, name)
if err != nil {
logger.Error("databricks_catalog_catalog.getCatalogCatalog", "api_error", err)
logger.Error("databricks_catalog.getCatalog", "api_error", err)
return nil, err
}

return *catalog, nil
}

func getCatalogCatalogWorkspaceBindings(ctx context.Context, d *plugin.QueryData, h *plugin.HydrateData) (interface{}, error) {
func getCatalogWorkspaceBindings(ctx context.Context, d *plugin.QueryData, h *plugin.HydrateData) (interface{}, error) {
logger := plugin.Logger(ctx)
name := h.Item.(catalog.CatalogInfo).Name

// Create client
client, err := connectDatabricksWorkspace(ctx, d)
if err != nil {
logger.Error("databricks_catalog_catalog.getCatalogCatalogWorkspaceBindings", "connection_error", err)
logger.Error("databricks_catalog.getCatalogWorkspaceBindings", "connection_error", err)
return nil, err
}

bindings, err := client.WorkspaceBindings.GetByName(ctx, name)
if err != nil {
logger.Error("databricks_catalog_catalog.getCatalogCatalogWorkspaceBindings", "api_error", err)
logger.Error("databricks_catalog.getCatalogWorkspaceBindings", "api_error", err)
return nil, err
}

Expand All @@ -240,13 +240,13 @@ func getCatalogPermissions(ctx context.Context, d *plugin.QueryData, h *plugin.H
// Create client
client, err := connectDatabricksWorkspace(ctx, d)
if err != nil {
logger.Error("databricks_catalog_catalog.getCatalogPermissions", "connection_error", err)
logger.Error("databricks_catalog.getCatalogPermissions", "connection_error", err)
return nil, err
}

permission, err := client.Grants.GetBySecurableTypeAndFullName(ctx, catalog.SecurableTypeCatalog, name)
if err != nil {
logger.Error("databricks_catalog_catalog.getCatalogPermissions", "api_error", err)
logger.Error("databricks_catalog.getCatalogPermissions", "api_error", err)
return nil, err
}
return permission.PrivilegeAssignments, nil
Expand All @@ -259,13 +259,13 @@ func getCatalogEffectivePermissions(ctx context.Context, d *plugin.QueryData, h
// Create client
client, err := connectDatabricksWorkspace(ctx, d)
if err != nil {
logger.Error("databricks_catalog_catalog.getCatalogEffectivePermissions", "connection_error", err)
logger.Error("databricks_catalog.getCatalogEffectivePermissions", "connection_error", err)
return nil, err
}

permission, err := client.Grants.GetEffectiveBySecurableTypeAndFullName(ctx, catalog.SecurableTypeCatalog, name)
if err != nil {
logger.Error("databricks_catalog_catalog.getCatalogEffectivePermissions", "api_error", err)
logger.Error("databricks_catalog.getCatalogEffectivePermissions", "api_error", err)
return nil, err
}
return permission.PrivilegeAssignments, nil
Expand Down
2 changes: 1 addition & 1 deletion databricks/table_databricks_catalog_schema.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ func tableDatabricksCatalogSchema(_ context.Context) *plugin.Table {
Name: "databricks_catalog_schema",
Description: "List schemas for a catalog in the metastore.",
List: &plugin.ListConfig{
ParentHydrate: listCatalogCatalogs,
ParentHydrate: listCatalogs,
Hydrate: listCatalogSchemas,
KeyColumns: plugin.OptionalColumns([]string{"catalog_name"}),
},
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,12 @@ import (

//// TABLE DEFINITION

func tableDatabricksWorkspaceWorkspace(_ context.Context) *plugin.Table {
func tableDatabricksWorkspace(_ context.Context) *plugin.Table {
return &plugin.Table{
Name: "databricks_workspace_workspace",
Name: "databricks_workspace",
Description: "List all secret workspaces available in the workspace.",
List: &plugin.ListConfig{
Hydrate: listWorkspaceWorkspaces,
Hydrate: listWorkspaces,
ShouldIgnoreError: isNotFoundError([]string{"RESOURCE_DOES_NOT_EXIST"}),
KeyColumns: plugin.OptionalColumns([]string{"path"}),
},
Expand Down Expand Up @@ -72,7 +72,7 @@ func tableDatabricksWorkspaceWorkspace(_ context.Context) *plugin.Table {

//// LIST FUNCTION

func listWorkspaceWorkspaces(ctx context.Context, d *plugin.QueryData, h *plugin.HydrateData) (interface{}, error) {
func listWorkspaces(ctx context.Context, d *plugin.QueryData, h *plugin.HydrateData) (interface{}, error) {
logger := plugin.Logger(ctx)
path := "/"

Expand All @@ -87,13 +87,13 @@ func listWorkspaceWorkspaces(ctx context.Context, d *plugin.QueryData, h *plugin
// Create client
client, err := connectDatabricksWorkspace(ctx, d)
if err != nil {
logger.Error("databricks_workspace_workspace.listWorkspaceWorkspaces", "connection_error", err)
logger.Error("databricks_workspace.listWorkspaces", "connection_error", err)
return nil, err
}

workspaces, err := client.Workspace.ListAll(ctx, request)
if err != nil {
logger.Error("databricks_workspace_workspace.listWorkspaceWorkspaces", "api_error", err)
logger.Error("databricks_workspace.listWorkspaces", "api_error", err)
return nil, err
}

Expand Down
32 changes: 18 additions & 14 deletions databricks/utils.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,14 @@ import (

func isNotFoundError(notFoundErrors []string) plugin.ErrorPredicate {
return func(err error) bool {
errMsg := err.(*apierr.APIError)
for _, msg := range notFoundErrors {
if strings.Contains(errMsg.ErrorCode, msg) {
return true
} else if strings.Contains(strconv.Itoa(errMsg.StatusCode), msg) {
return true
switch err := err.(type) {
case *apierr.APIError:
for _, msg := range notFoundErrors {
if strings.Contains(err.ErrorCode, msg) {
return true
} else if strings.Contains(strconv.Itoa(err.StatusCode), msg) {
return true
}
}
}
return false
Expand All @@ -25,14 +27,16 @@ func isNotFoundError(notFoundErrors []string) plugin.ErrorPredicate {

func shouldRetryError(retryErrors []string) plugin.ErrorPredicateWithContext {
return func(ctx context.Context, d *plugin.QueryData, h *plugin.HydrateData, err error) bool {
errMsg := err.(*apierr.APIError)
for _, msg := range retryErrors {
if strings.Contains(errMsg.ErrorCode, msg) {
plugin.Logger(ctx).Error("databricks_errors.shouldRetryError", "rate_limit_error", err)
return true
} else if strings.Contains(strconv.Itoa(errMsg.StatusCode), msg) {
plugin.Logger(ctx).Error("databricks_errors.shouldRetryError", "rate_limit_error", err)
return true
switch err := err.(type) {
case *apierr.APIError:
for _, msg := range retryErrors {
if strings.Contains(err.ErrorCode, msg) {
plugin.Logger(ctx).Error("databricks_errors.shouldRetryError", "rate_limit_error", err)
return true
} else if strings.Contains(strconv.Itoa(err.StatusCode), msg) {
plugin.Logger(ctx).Error("databricks_errors.shouldRetryError", "rate_limit_error", err)
return true
}
}
}
return false
Expand Down
20 changes: 10 additions & 10 deletions docs/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -114,27 +114,27 @@ connection "databricks" {
}
```

By default, all options are commented out in the default connection, thus Steampipe will resolve your credentials using the same mechanism as the Databricks CLI (Databricks environment variables, default profile, etc). This provides a quick way to get started with Steampipe, but you will probably want to customize your experience using configuration options for [querying multiple regions](#multi-account-connections), [configuring credentials](#configuring-databricks-credentials) from your [Databricks Profiles](#databricks-profile-credentials).
You can customize your experience using configuration options for [querying multiple accounts](#multi-account-connections), [configuring credentials](#configuring-databricks-credentials) from your [Databricks Profiles](#databricks-profile-credentials).
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we can add support to the plugin to pickup the DEFAULT profile by default if no other config args for auth are mentioned, we can re-add this line in with the appropriate info. I'm good with the change in this PR though as is.


## Multi-Account Connections

You may create multiple databricks connections:
```hcl
connection "databricks_dev" {
plugin = "databricks"
profile = "databricks_dev"
config_profile = "databricks_dev"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can this config arg just be called profile instead? Are there multiple types of profiles you can use with the Databricks CLI/SDK?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reference points for naming:

For env vars:

  • DATABRICKS_CONFIG_PROFILE
  • DATABRICKS_CONFIG_FILE

From Terraform (https://registry.terraform.io/providers/databricks/databricks/latest/docs#argument-reference):

  • profile
  • config_file

In their Config struct (https://github.com/databricks/databricks-sdk-go/blob/main/config/config.go#L54-L60):

  • Profile
  • ConfigFile

account_id = abcdd0f81-9be0-4425-9e29-3a7d96782373
}

connection "databricks_qa" {
plugin = "databricks"
profile = "databricks_qa"
config_profile = "databricks_qa"
account_id = wxyzd0f81-9be0-4425-9e29-3a7d96782373
}

connection "databricks_prod" {
plugin = "databricks"
profile = "databricks_prod"
config_profile = "databricks_prod"
account_id = pqrsd0f81-9be0-4425-9e29-3a7d96782373
}
```
Expand Down Expand Up @@ -206,19 +206,19 @@ account_id = abcdd0f81-9be0-4425-9e29-3a7d96782373
```hcl
connection "databricks_user1-account" {
plugin = "databricks"
profile = "user1-account"
config_profile = "user1-account"
account_id = "abcdd0f81-9be0-4425-9e29-3a7d96782373"
}

connection "databricks_user1-workspace" {
plugin = "databricks"
profile = "user1-workspace"
config_profile = "user1-workspace"
account_id = "abcdd0f81-9be0-4425-9e29-3a7d96782373"
}

connection "databricks_user1-basic" {
plugin = "databricks"
profile = "user1-basic"
config_profile = "user1-basic"
account_id = "abcdd0f81-9be0-4425-9e29-3a7d96782373"
}
```
Expand All @@ -241,7 +241,7 @@ account_id = abcdd0f81-9be0-4425-9e29-3a7d96782373
```hcl
connection "databricks_user1-account" {
plugin = "databricks"
profile = "user1-account"
config_profile = "user1-account"
account_id = "abcdd0f81-9be0-4425-9e29-3a7d96782373"
}
```
Expand All @@ -264,7 +264,7 @@ account_id = abcdd0f81-9be0-4425-9e29-3a7d96782373
```hcl
connection "databricks_user1-workspace" {
plugin = "databricks"
profile = "user1-workspace"
config_profile = "user1-workspace"
account_id = "abcdd0f81-9be0-4425-9e29-3a7d96782373"
}
```
Expand All @@ -291,7 +291,7 @@ connection "databricks_user1-workspace" {

### Credentials from Environment Variables

Alternatively, you can also use the standard Databricks environment variables to obtain credentials **only if other argument (`profile`, `account_id`, `account_token`/`account_host`/`workspace_token`/`workspace_host`) is not specified** in the connection:
Alternatively, you can also use the standard Databricks environment variables to obtain credentials **only if other argument (`config_profile`, `account_id`, `account_token`/`account_host`/`workspace_token`/`workspace_host`) is not specified** in the connection:

```sh
export DATABRICKS_CONFIG_PROFILE=user1-test
Expand Down
Loading