steve/databricks-cli
steve
/
databricks-cli
mirror of https://github.com/databricks/cli.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Comments

This commit is contained in:
Pieter Noordhuis 2024-11-13 13:40:41 +01:00
parent 1ab0a59ba5
commit b42fbe35a9
No known key found for this signature in database
GPG Key ID: 12ACCCC104CF2930
1 changed files with 11 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
libs/dbr/context.go
View File

@ -2,7 +2,10 @@ package dbr
import "context" import "context"
// key is a private type to prevent collisions with other packages. // key is a package-local type to use for context keys.
//
// Using an unexported type for context keys prevents key collisions across
// packages since external packages cannot create values of this type.
type key int type key int
const ( const (
@ -14,7 +17,7 @@ const (
// DetectRuntime detects whether or not the current // DetectRuntime detects whether or not the current
// process is running inside a Databricks Runtime environment. // process is running inside a Databricks Runtime environment.
// It return a new context with the detection result cached. // It return a new context with the detection result set.
func DetectRuntime(ctx context.Context) context.Context { func DetectRuntime(ctx context.Context) context.Context {
if v := ctx.Value(dbrKey); v != nil { if v := ctx.Value(dbrKey); v != nil {
panic("dbr.DetectRuntime called twice on the same context") panic("dbr.DetectRuntime called twice on the same context")
@ -23,7 +26,7 @@ func DetectRuntime(ctx context.Context) context.Context {
} }
// MockRuntime is a helper function to mock the detection result. // MockRuntime is a helper function to mock the detection result.
// It returns a new context with the detection result cached. // It returns a new context with the detection result set.
func MockRuntime(ctx context.Context, b bool) context.Context { func MockRuntime(ctx context.Context, b bool) context.Context {
if v := ctx.Value(dbrKey); v != nil { if v := ctx.Value(dbrKey); v != nil {
panic("dbr.MockRuntime called twice on the same context") panic("dbr.MockRuntime called twice on the same context")
@ -31,8 +34,12 @@ func MockRuntime(ctx context.Context, b bool) context.Context {
return context.WithValue(ctx, dbrKey, b) return context.WithValue(ctx, dbrKey, b)
} }
// RunsOnRuntime returns the cached detection result from the context. // RunsOnRuntime returns the detection result from the context.
// It expects a context returned by [DetectRuntime] or [MockRuntime]. // It expects a context returned by [DetectRuntime] or [MockRuntime].
//
// We store this value in a context to avoid having to use either
// a global variable, passing a boolean around everywhere, or
// performing the same detection multiple times.
func RunsOnRuntime(ctx context.Context) bool { func RunsOnRuntime(ctx context.Context) bool {
v := ctx.Value(dbrKey) v := ctx.Value(dbrKey)
if v == nil { if v == nil {