99 dbsql "github.com/databricks/databricks-sql-go"
1010 "github.com/databricks/databricks-sql-go/auth"
1111 "github.com/databricks/databricks-sql-go/auth/oauth/m2m"
12- "github.com/grafana/grafana-azure-sdk-go/azsettings"
1312 "github.com/grafana/grafana-plugin-sdk-go/backend"
1413 "github.com/grafana/grafana-plugin-sdk-go/backend/instancemgmt"
1514 "github.com/grafana/grafana-plugin-sdk-go/backend/log"
@@ -70,7 +69,6 @@ type ConnectionSettings struct {
7069 MaxRetryDuration time.Duration
7170 Timeout time.Duration
7271 MaxRows int
73- idToken string
7472}
7573
7674// NewSampleDatasource creates a new datasource instance.
@@ -93,7 +91,7 @@ func NewSampleDatasource(ctx context.Context, settings backend.DataSourceInstanc
9391 port = portInt
9492 }
9593
96- if datasourceSettings .AuthenticationMethod == "m2m" || datasourceSettings .AuthenticationMethod == "oauth2_client_credentials" || datasourceSettings . AuthenticationMethod == "azure_ad_forward" {
94+ if datasourceSettings .AuthenticationMethod == "m2m" || datasourceSettings .AuthenticationMethod == "oauth2_client_credentials" {
9795 var authenticator auth.Authenticator
9896
9997 if datasourceSettings .AuthenticationMethod == "oauth2_client_credentials" {
@@ -114,15 +112,6 @@ func NewSampleDatasource(ctx context.Context, settings backend.DataSourceInstanc
114112 datasourceSettings .Hostname ,
115113 []string {},
116114 )
117- } else if datasourceSettings .AuthenticationMethod == "azure_ad_forward" {
118- azureSettings , err := azsettings .ReadSettings (ctx )
119- if err != nil {
120- log .DefaultLogger .Info ("Failed to get Azure Setting" , "err" , err )
121- return nil , err
122- }
123- authenticator = integrations .NewAzureADCredentials (
124- azureSettings ,
125- )
126115 } else {
127116 log .DefaultLogger .Info ("Authentication Method Parse Error" , "err" , nil )
128117 return nil , fmt .Errorf ("authentication Method Parse Error" )
@@ -142,12 +131,20 @@ func NewSampleDatasource(ctx context.Context, settings backend.DataSourceInstanc
142131 return nil , err
143132 } else {
144133 log .DefaultLogger .Info ("Init Databricks SQL DB" )
134+ databricksDB := sql .OpenDB (connector )
135+
136+ if err := databricksDB .Ping (); err != nil {
137+ log .DefaultLogger .Info ("Ping Error (Could not ping Databricks)" , "err" , err )
138+ return nil , err
139+ }
145140
141+ SetDatasourceSettings (databricksDB , connectionSettings )
146142 log .DefaultLogger .Info ("Store Databricks SQL DB Connection" )
147143 return & Datasource {
148144 connector : connector ,
149- databricksDB : nil ,
145+ databricksDB : databricksDB ,
150146 connectionSettings : connectionSettings ,
147+ datasourceSettings : * datasourceSettings ,
151148 }, nil
152149 }
153150 } else if datasourceSettings .AuthenticationMethod == "dsn" || datasourceSettings .AuthenticationMethod == "" {
@@ -179,8 +176,17 @@ func NewSampleDatasource(ctx context.Context, settings backend.DataSourceInstanc
179176 connector : connector ,
180177 databricksDB : databricksDB ,
181178 connectionSettings : connectionSettings ,
179+ datasourceSettings : * datasourceSettings ,
182180 }, nil
183181
182+ } else if datasourceSettings .AuthenticationMethod == "azure_entra_pass_thru" {
183+
184+ return & Datasource {
185+ connector : nil ,
186+ databricksDB : nil ,
187+ connectionSettings : connectionSettings ,
188+ datasourceSettings : * datasourceSettings ,
189+ }, nil
184190 }
185191
186192 return nil , fmt .Errorf ("Invalid Connection Method" )
@@ -302,6 +308,8 @@ type Datasource struct {
302308 connector driver.Connector
303309 databricksDB * sql.DB
304310 connectionSettings ConnectionSettings
311+ datasourceSettings DatasourceSettings
312+ token string
305313}
306314
307315func (d * Datasource ) CallResource (ctx context.Context , req * backend.CallResourceRequest , sender backend.CallResourceResponseSender ) error {
@@ -322,6 +330,15 @@ func (d *Datasource) Dispose() {
322330func (d * Datasource ) QueryData (ctx context.Context , req * backend.QueryDataRequest ) (* backend.QueryDataResponse , error ) {
323331 log .DefaultLogger .Info ("QueryData called" , "request" , req )
324332
333+ if d .datasourceSettings .AuthenticationMethod == "azure_entra_pass_thru" {
334+ token := req .GetHTTPHeader (backend .OAuthIdentityTokenHeaderName )
335+ err := d .CheckAzureEntraPassThru (token )
336+ if err != nil {
337+ log .DefaultLogger .Error ("Azure Entra Connection Failed" , "err" , err )
338+ return nil , err
339+ }
340+ }
341+
325342 // create response struct
326343 response := backend .NewQueryDataResponse ()
327344
@@ -446,28 +463,74 @@ func (d *Datasource) query(ctx context.Context, pCtx backend.PluginContext, quer
446463 return response
447464}
448465
466+ func (d * Datasource ) CheckAzureEntraPassThru (token string ) error {
467+ if token == "" {
468+ log .DefaultLogger .Info ("Token is empty" )
469+ return fmt .Errorf ("no Azure Entra Token provided" )
470+ }
471+ if d .databricksDB != nil && token == d .token {
472+ return nil
473+ }
474+ if token != d .token {
475+ log .DefaultLogger .Info ("Token is different" )
476+ d .token = strings .TrimPrefix (token , "Bearer " )
477+ }
478+ port := 443
479+ if d .datasourceSettings .Port != "" {
480+ portInt , err := strconv .Atoi (d .datasourceSettings .Port )
481+ if err != nil {
482+ log .DefaultLogger .Info ("Port Parse Error" , "err" , err )
483+ return err
484+ }
485+ port = portInt
486+ }
487+
488+ connector , err := dbsql .NewConnector (
489+ dbsql .WithAccessToken (d .token ),
490+ dbsql .WithServerHostname (d .datasourceSettings .Hostname ),
491+ dbsql .WithHTTPPath (d .datasourceSettings .Path ),
492+ dbsql .WithPort (port ),
493+ dbsql .WithTimeout (d .connectionSettings .Timeout ),
494+ dbsql .WithMaxRows (d .connectionSettings .MaxRows ),
495+ dbsql .WithRetries (d .connectionSettings .Retries , d .connectionSettings .RetryBackoff , d .connectionSettings .MaxRetryDuration ),
496+ )
497+ if err != nil {
498+ log .DefaultLogger .Info ("Connector Error" , "err" , err )
499+ return err
500+ } else {
501+ log .DefaultLogger .Info ("Init Databricks SQL DB" )
502+ databricksDB := sql .OpenDB (connector )
503+
504+ if err := databricksDB .Ping (); err != nil {
505+ log .DefaultLogger .Info ("Ping Error (Could not ping Databricks)" , "err" , err )
506+ return err
507+ }
508+
509+ SetDatasourceSettings (databricksDB , d .connectionSettings )
510+ log .DefaultLogger .Info ("Store Databricks SQL DB Connection" )
511+ d .connector = connector
512+ d .databricksDB = databricksDB
513+ }
514+ return nil
515+ }
516+
449517// CheckHealth handles health checks sent from Grafana to the plugin.
450518// The main use case for these health checks is the test button on the
451519// datasource configuration page which allows users to verify that
452520// a datasource is working as expected.
453521func (d * Datasource ) CheckHealth (ctx context.Context , req * backend.CheckHealthRequest ) (* backend.CheckHealthResult , error ) {
454- log . DefaultLogger . Info ( "CheckHealth called" , "request" , req )
522+ if d . datasourceSettings . AuthenticationMethod == "azure_entra_pass_thru" {
455523
456- token := strings .Fields (req .GetHTTPHeader (backend .OAuthIdentityTokenHeaderName ))
457- idToken := req .GetHTTPHeader (backend .OAuthIdentityIDTokenHeaderName )
458- log .DefaultLogger .Info ("Token" , "token" , token )
459- log .DefaultLogger .Info ("ID Token" , "idToken" , idToken )
460-
461- ctx = context .WithValue (ctx , backend .OAuthIdentityTokenHeaderName , token )
462- ctx = context .WithValue (ctx , backend .OAuthIdentityIDTokenHeaderName , idToken )
463-
464- if d .databricksDB == nil {
465- err := d .RefreshDBConnection ()
524+ token := req .GetHTTPHeader (backend .OAuthIdentityTokenHeaderName )
525+ err := d .CheckAzureEntraPassThru (token )
466526 if err != nil {
467- log .DefaultLogger .Info ("RefreshDBConnection Error" , "err" , err )
468- return nil , err
527+ return & backend.CheckHealthResult {
528+ Status : backend .HealthStatusError ,
529+ Message : fmt .Sprintf ("Azure Entra Connection Failed: %s" , err ),
530+ }, nil
469531 }
470532 }
533+ log .DefaultLogger .Info ("CheckHealth called" , "request" , req )
471534
472535 rows , err := d .QueryContext (ctx , "SELECT 1" )
473536
0 commit comments