Add Databricks Support

README.md

@@ -15,6 +15,7 @@ Currently supported:
 - Postgres
 - ClickHouse
 - AWS Athena
+- Databricks
 - MS-SQL
 - MySQL
 - Snowflake
@@ -229,6 +230,7 @@ For some database backends some special functionality is available:
 
 * cloudsql-postgres: A special `*` character can be used to query all databases accessible by the account
 * cloudsql-mysql: Same as above
+* databricks: Supports access token, OAuth M2M (machine-to-machine) or U2M (user-to-machine) authentication as supported by the [driver](https://github.com/databricks/databricks-sql-go).
 * rds-postgres: This type of URL expects a working AWS configuration
   which will use the equivalent of `rds generate-db-auth-token`
   for the password. For this driver, the `AWS_REGION` environment variable

config.go

@@ -165,8 +165,15 @@ type connection struct {
 	user                string
 	tokenExpirationTime time.Time
 	iteratorValues      []string
-    	snowflakeConfig *gosnowflake.Config
-    	snowflakeDSN    string
+	snowflakeConfig     *gosnowflake.Config
+	snowflakeDSN        string
+	// Databricks OAuth configuration
+	databricksOAuthType    string // "m2m" or "u2m"
+	databricksClientID     string
+	databricksClientSecret string
+	databricksHost         string
+	databricksHTTPPath     string
+	databricksCatalog      string
 }
 
 // Query is an SQL query that is executed on a connection

config.yml.dist

@@ -128,7 +128,7 @@ jobs:
   interval: '5m'
   connections:
   # see https://godoc.org/github.com/segmentio/go-athena#Driver.Open
-  - 'athena://HOST_VALUE_IGNORED/<DB_NAME>?db=<DB_NAME>&region=<AWS_REGION>&output_location=s3://aws-athena-query-results-<ACCOUNT_ID>-<REGION>'
+  - 'athena://?db=<DB_NAME>&region=<AWS_REGION>&output_location=s3://aws-athena-query-results-<ACCOUNT_ID>-<REGION>'
   queries:
   - name: "athena_query_rows"
     help: "Number of rows..."
@@ -193,3 +193,25 @@ jobs:
     query:  |
             SELECT datname::text, usename::text, COUNT(*)::float AS count
             FROM pg_stat_activity GROUP BY datname, usename;
+- name: "databricks"
+  interval: '5m'
+  connections:
+  # Connect to Databricks SQL Warehouse or All Purpose Cluster
+  # Connect with M2M
+  # - databricks://your-workspace.cloud.databricks.com/sql/1.0/warehouses/abc123?catalog=main&authType=m2m&clientId=YOUR_CLIENT_ID&clientSecret=YOUR_CLIENT_SECRET
+  # Connect with a token
+  - 'databricks://token:TOKEN@INSTANCE.cloud.databricks.com:443/sql/1.0/warehouses/WAREHOUSE_ID'
+  queries:
+  - name: "databricks_usage"
+    help: "Total databricks usage by date"
+    values:
+      - "total"
+    query: |
+        SELECT 
+            usage_date,
+            SUM(usage_quantity) AS total
+        FROM 
+            system.billing.usage
+        WHERE 
+            usage_date = CURRENT_DATE()
+        GROUP BY usage_date;

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/ClickHouse/clickhouse-go/v2 v2.34.0
 	github.com/aws/aws-sdk-go v1.50.6
 	github.com/cenkalti/backoff v2.2.1+incompatible
+	github.com/databricks/databricks-sql-go v1.6.1
 	github.com/go-kit/log v0.2.1
 	github.com/go-sql-driver/mysql v1.8.1
 	github.com/gobwas/glob v0.2.3
@@ -39,6 +40,8 @@ require (
 	github.com/JohnCGriffin/overflow v0.0.0-20211019200055-46fa312c352c // indirect
 	github.com/andybalholm/brotli v1.1.1 // indirect
 	github.com/apache/arrow-go/v18 v18.0.0 // indirect
+	github.com/apache/arrow/go/v12 v12.0.1 // indirect
+	github.com/apache/thrift v0.21.0 // indirect
 	github.com/aws/aws-sdk-go-v2 v1.26.1 // indirect
 	github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.2 // indirect
 	github.com/aws/aws-sdk-go-v2/credentials v1.17.11 // indirect
@@ -54,14 +57,19 @@ require (
 	github.com/aws/smithy-go v1.20.2 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
+	github.com/coreos/go-oidc/v3 v3.5.0 // indirect
 	github.com/danieljoos/wincred v1.2.2 // indirect
+	github.com/dnephin/pflag v1.0.7 // indirect
 	github.com/dvsekhvalnov/jose2go v1.6.0 // indirect
 	github.com/elastic/go-sysinfo v1.11.2 // indirect
 	github.com/elastic/go-windows v1.0.1 // indirect
+	github.com/fatih/color v1.15.0 // indirect
 	github.com/felixge/httpsnoop v1.0.4 // indirect
+	github.com/fsnotify/fsnotify v1.5.4 // indirect
 	github.com/gabriel-vasile/mimetype v1.4.7 // indirect
 	github.com/go-faster/city v1.0.1 // indirect
 	github.com/go-faster/errors v0.7.1 // indirect
+	github.com/go-jose/go-jose/v3 v3.0.0 // indirect
 	github.com/go-logfmt/logfmt v0.6.0 // indirect
 	github.com/go-logr/logr v1.4.2 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect
@@ -71,12 +79,16 @@ require (
 	github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9 // indirect
 	github.com/golang-sql/sqlexp v0.1.0 // indirect
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
+	github.com/golang/snappy v0.0.4 // indirect
 	github.com/google/flatbuffers v24.12.23+incompatible // indirect
 	github.com/google/s2a-go v0.1.8 // indirect
+	github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 // indirect
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/googleapis/enterprise-certificate-proxy v0.3.4 // indirect
 	github.com/googleapis/gax-go/v2 v2.13.0 // indirect
 	github.com/gsterjov/go-libsecret v0.0.0-20161001094733-a6f4afe4910c // indirect
+	github.com/hashicorp/go-cleanhttp v0.5.1 // indirect
+	github.com/hashicorp/go-retryablehttp v0.7.1 // indirect
 	github.com/hashicorp/go-uuid v1.0.3 // indirect
 	github.com/jackc/chunkreader/v2 v2.0.1 // indirect
 	github.com/jackc/pgconn v1.14.3 // indirect
@@ -94,8 +106,13 @@ require (
 	github.com/jcmturner/rpc/v2 v2.0.3 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
 	github.com/joeshaw/multierror v0.0.0-20140124173710-69b34d4ec901 // indirect
+	github.com/klauspost/asmfmt v1.3.2 // indirect
 	github.com/klauspost/compress v1.17.11 // indirect
 	github.com/klauspost/cpuid/v2 v2.2.9 // indirect
+	github.com/mattn/go-colorable v0.1.13 // indirect
+	github.com/mattn/go-isatty v0.0.19 // indirect
+	github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 // indirect
+	github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 // indirect
 	github.com/mtibben/percent v0.2.1 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/paulmach/orb v0.11.1 // indirect
@@ -104,6 +121,7 @@ require (
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/prometheus/client_model v0.6.1 // indirect
 	github.com/prometheus/procfs v0.15.1 // indirect
+	github.com/rs/zerolog v1.28.0 // indirect
 	github.com/segmentio/asm v1.2.0 // indirect
 	github.com/shopspring/decimal v1.4.0 // indirect
 	github.com/sirupsen/logrus v1.9.3 // indirect
@@ -130,5 +148,6 @@ require (
 	google.golang.org/grpc v1.67.1 // indirect
 	google.golang.org/protobuf v1.35.1 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
+	gotest.tools/gotestsum v1.8.2 // indirect
 	howett.net/plist v1.0.1 // indirect
 )