🐛 Fix: use Invalid Session Handle retry logic for Code Auto Complete …

…Queries (#83)

CHANGELOG.md

@@ -1,6 +1,12 @@
 # Changelog
 
-## 1.2.7
+## 1.2.8
+
+- Bugfix: Use retry logic on Invalid SessionHandle error for Code Autocomplete Suggestions Queries
+
+---
+
+### 1.2.7
 
 - Feature: Add support to fetch Databricks Token from external OAuth Endpoint 
 

package.json

@@ -1,7 +1,7 @@
 {
   "name": "mullerpeter-databricks-datasource",
   "private": true,
-  "version": "1.2.7",
+  "version": "1.2.8",
   "description": "Databricks SQL Connector",
   "scripts": {
     "build": "webpack -c ./.config/webpack/webpack.config.ts --env production",

pkg/plugin/helper.go

@@ -25,7 +25,7 @@ type defaultsResponseBody struct {
 	DefaultSchema  string `json:"defaultSchema"`
 }
 
-func autocompletionQueries(req *backend.CallResourceRequest, sender backend.CallResourceResponseSender, db *sql.DB) error {
+func autocompletionQueries(req *backend.CallResourceRequest, sender backend.CallResourceResponseSender, d *Datasource) error {
 	path := req.Path
 	log.DefaultLogger.Info("CallResource called", "path", path)
 	var body schemaRequestBody
@@ -36,7 +36,7 @@ func autocompletionQueries(req *backend.CallResourceRequest, sender backend.Call
 	}
 	switch path {
 	case "catalogs":
-		rows, err := db.Query("SHOW CATALOGS")
+		rows, err := d.ExecuteQuery("SHOW CATALOGS")
 		if err != nil {
 			log.DefaultLogger.Error("CallResource Error", "err", err)
 			return err
@@ -74,7 +74,7 @@ func autocompletionQueries(req *backend.CallResourceRequest, sender backend.Call
 			queryString = fmt.Sprintf("SHOW SCHEMAS IN %s", body.Catalog)
 		}
 		log.DefaultLogger.Info("CallResource called", "queryString", queryString)
-		rows, err := db.Query(queryString)
+		rows, err := d.ExecuteQuery(queryString)
 		if err != nil {
 			log.DefaultLogger.Error("CallResource Error", "err", err)
 			return err
@@ -114,7 +114,7 @@ func autocompletionQueries(req *backend.CallResourceRequest, sender backend.Call
 			}
 		}
 		log.DefaultLogger.Info("CallResource called", "queryString", queryString)
-		rows, err := db.Query(queryString)
+		rows, err := d.ExecuteQuery(queryString)
 		if err != nil {
 			log.DefaultLogger.Error("CallResource Error", "err", err)
 			return err
@@ -150,7 +150,7 @@ func autocompletionQueries(req *backend.CallResourceRequest, sender backend.Call
 	case "columns":
 		queryString := fmt.Sprintf("DESCRIBE TABLE %s", body.Table)
 		log.DefaultLogger.Info("CallResource called", "queryString", queryString)
-		rows, err := db.Query(queryString)
+		rows, err := d.ExecuteQuery(queryString)
 		if err != nil {
 			log.DefaultLogger.Error("CallResource Error", "err", err)
 			return err
@@ -190,11 +190,19 @@ func autocompletionQueries(req *backend.CallResourceRequest, sender backend.Call
 	case "defaults":
 		queryString := "SELECT current_catalog(), current_schema();"
 		log.DefaultLogger.Info("CallResource called", "queryString", queryString)
-		row := db.QueryRow(queryString)
+		rows, err := d.ExecuteQuery(queryString)
+		if err != nil {
+			log.DefaultLogger.Error("CallResource Error", "err", err)
+			return err
+		}
+		defer rows.Close()
 		var currentCatalog sql.NullString
 		var currentSchema sql.NullString
-
-		err := row.Scan(&currentCatalog, &currentSchema)
+		if rows.Next() == false {
+			log.DefaultLogger.Error("CallResource Error", "err", "No rows returned")
+			return fmt.Errorf("no rows returned")
+		}
+		err = rows.Scan(&currentCatalog, &currentSchema)
 		if err != nil {
 			log.DefaultLogger.Error("CallResource Error", "err", err)
 			return err

pkg/plugin/plugin.go

@@ -187,7 +187,7 @@ type Datasource struct {
 }
 
 func (d *Datasource) CallResource(ctx context.Context, req *backend.CallResourceRequest, sender backend.CallResourceResponseSender) error {
-	return autocompletionQueries(req, sender, d.databricksDB)
+	return autocompletionQueries(req, sender, d)
 }
 
 // Dispose here tells plugin SDK that plugin wants to clean up resources when a new instance