Introducing some basic REST API calls for SQL Warehouses

Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "databricks"
+version = "0.1.0"
+edition = "2021"
+license = "LICENSE"
+readme = "README.md"
+repository = "https://github.com/buoyant-data/databricks"
+
+[dependencies]
+feignhttp = { version = "0.5.1", features = ["json"]}
+once_cell = "1"
+serde = { version = "1", features = ["derive"]}
+url = { version = "2", features = ["serde"] }
+
+[dev-dependencies]
+anyhow = "1"
+tokio = { version = "1", features = ["full"]}

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 R Tyler Croy
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

examples/create_warehouse.rs

@@ -0,0 +1,20 @@
+use databricks::sql::Warehouse;
+
+#[tokio::main]
+async fn main() -> Result<(), anyhow::Error> {
+    let config = databricks::Config::default();
+
+    println!("Starting example..");
+    let warehouse = Warehouse::builder()
+                        .name("rust-test")
+                        .cluster_size("Large")
+
+
+    for warehouse in Warehouse::list(&config).await? {
+        println!(
+            "Warehouse {} is {} and has {} clusters",
+            warehouse.name, warehouse.state, warehouse.num_clusters
+        );
+    }
+    Ok(())
+}

examples/list_warehouses.rs

@@ -0,0 +1,16 @@
+use databricks::sql::Warehouse;
+
+#[tokio::main]
+async fn main() -> Result<(), anyhow::Error> {
+    let config = databricks::Config::default();
+
+    println!("Starting example..");
+    for warehouse in Warehouse::list(&config).await? {
+        println!(
+            "Warehouse {} is {} and has {} clusters",
+            warehouse.name, warehouse.state, warehouse.num_clusters
+        );
+        //println!("{warehouse:?}");
+    }
+    Ok(())
+}

src/lib.rs

@@ -0,0 +1,32 @@
+use once_cell::sync::Lazy;
+
+pub mod sql;
+
+static BASE_URL: Lazy<String> = Lazy::new(|| {
+    std::env::var("DATABRICKS_URL").unwrap_or("https://cloud.databricks.com".to_string())
+});
+
+/// Genericc configuration
+#[derive(Clone, Debug)]
+pub struct Config {
+    token: String,
+}
+
+impl Config {
+    pub fn with_token<T: Into<String>>(token: T) -> Self {
+        Self {
+            token: token.into(),
+        }
+    }
+}
+
+impl Default for Config {
+    fn default() -> Self {
+        Self::with_token(
+            std::env::var("DATABRICKS_TOKEN")
+            .expect("Need DATABRICKS_TOKEN set in the environment"))
+    }
+}
+
+#[cfg(test)]
+mod tests {}

src/sql.rs

@@ -0,0 +1,214 @@
+//! Databricks SQL REST APIs
+
+use serde::Deserialize;
+use url::Url;
+
+use std::collections::HashMap;
+
+/// Rust representation of a Warehouse structure
+#[allow(unused)]
+#[derive(Clone, Debug, Deserialize)]
+pub struct Warehouse {
+    #[serde(default, skip_deserializing)]
+    config: crate::Config,
+    pub id: String,
+    pub name: String,
+    pub state: String,
+    cluster_size: String,
+    min_num_clusters: u64,
+    max_num_clusters: u64,
+    pub num_clusters: u64,
+    num_active_sessions: u64,
+    spot_instance_policy: String,
+    enable_photon: bool,
+    auto_stop_mins: u64,
+    creator_name: String,
+    instance_profile_arn: String,
+    warehouse_type: String,
+    odbc_params: ODBCParameters,
+    jdbc_url: Url,
+    health: Option<Health>,
+    channel: Option<Channel>,
+}
+
+impl Warehouse {
+    /// Start the warehouse
+    pub async fn start(&self) -> feignhttp::Result<String> {
+        api::start_warehouse(&self.id, &self.config.token).await
+    }
+
+    /// Stop the warehouse
+    pub async fn stop(&self) -> feignhttp::Result<String> {
+        api::stop_warehouse(&self.id, &self.config.token).await
+    }
+
+    /// Create a new Warehouse
+    ///
+    /// ```
+    /// {
+    ///   "name": "string",
+    ///   "cluster_size": "string",
+    ///   "min_num_clusters": "1",
+    ///   "max_num_clusters": 0,
+    ///   "auto_stop_mins": "120",
+    ///   "creator_name": "string",
+    ///   "instance_profile_arn": "string",
+    ///   "tags": {
+    ///     "custom_tags": [
+    ///       {
+    ///         "key": "string",
+    ///         "value": "string"
+    ///       }
+    ///     ]
+    ///   },
+    ///   "spot_instance_policy": "POLICY_UNSPECIFIED",
+    ///   "enable_photon": true,
+    ///   "channel": {
+    ///     "name": "CHANNEL_NAME_UNSPECIFIED",
+    ///     "dbsql_version": "string"
+    ///   },
+    ///   "enable_serverless_compute": true,
+    ///   "warehouse_type": "TYPE_UNSPECIFIED"
+    /// }
+    /// ```
+    pub async fn create(&self, config: &crate::Config) -> feignhttp::Result<String> {
+
+    }
+
+    /// List the warehouses visible to the configured user
+    pub async fn list(config: &crate::Config) -> feignhttp::Result<Vec<Warehouse>> {
+        let mut response = api::warehouses(&config.token).await?;
+        Ok(response
+            .warehouses
+            .drain(..)
+            .map(|mut w| { w.config = config.clone(); w })
+            .collect())
+    }
+}
+
+/// Rust representation of a warehouse's channel
+#[allow(unused)]
+#[derive(Clone, Debug, Deserialize)]
+pub struct Channel {
+    name: Option<String>,
+    dbsql_version: Option<String>,
+}
+
+/// Rust representation of a warehouse's ODBC parameters
+#[allow(unused)]
+#[derive(Clone, Debug, Deserialize)]
+pub struct ODBCParameters {
+    hostname: String,
+    path: String,
+    protocol: String,
+    port: u64,
+}
+
+/// Rust representation of a warehouse's health
+#[allow(unused)]
+#[derive(Clone, Debug, Deserialize)]
+pub struct Health {
+    status: String,
+    message: Option<String>,
+    summary: Option<String>,
+    details: Option<String>,
+    failure_reason: Option<FailureReason>,
+}
+
+#[allow(unused)]
+#[derive(Clone, Debug, Deserialize)]
+pub struct FailureReason {
+    code: String,
+    r#type: String,
+    parameters: HashMap<String, String>,
+}
+
+mod api {
+    use serde::Deserialize;
+    use feignhttp::{get, post};
+
+    /// List the SQL Warehouses accessible to the user
+    ///
+    /// [More documentation here](https://docs.databricks.com/api/workspace/warehouses/list)
+    ///
+    /// ```json
+    /// {
+    ///   "warehouses": [
+    ///     {
+    ///       "id": "string",
+    ///       "name": "string",
+    ///       "cluster_size": "string",
+    ///       "min_num_clusters": "1",
+    ///       "max_num_clusters": 0,
+    ///       "auto_stop_mins": "120",
+    ///       "creator_name": "string",
+    ///       "instance_profile_arn": "string",
+    ///       "tags": {
+    ///         "custom_tags": [
+    ///           {
+    ///             "key": "string",
+    ///             "value": "string"
+    ///           }
+    ///         ]
+    ///       },
+    ///       "spot_instance_policy": "POLICY_UNSPECIFIED",
+    ///       "enable_photon": true,
+    ///       "channel": {
+    ///         "name": "CHANNEL_NAME_UNSPECIFIED",
+    ///         "dbsql_version": "string"
+    ///       },
+    ///       "enable_serverless_compute": true,
+    ///       "warehouse_type": "TYPE_UNSPECIFIED",
+    ///       "num_clusters": 0,
+    ///       "num_active_sessions": 0,
+    ///       "state": "STARTING",
+    ///       "jdbc_url": "string",
+    ///       "odbc_params": {
+    ///         "hostname": "string",
+    ///         "path": "string",
+    ///         "protocol": "string",
+    ///         "port": 0
+    ///       },
+    ///       "health": {
+    ///         "status": "STATUS_UNSPECIFIED",
+    ///         "message": "string",
+    ///         "failure_reason": {
+    ///           "code": "UNKNOWN",
+    ///           "type": "SUCCESS",
+    ///           "parameters": {
+    ///             "property1": "string",
+    ///             "property2": "string"
+    ///           }
+    ///         },
+    ///         "summary": "string",
+    ///         "details": "string"
+    ///       }
+    ///     }
+    ///   ]
+    /// }
+    /// ```
+    #[get(url = crate::BASE_URL,
+        headers = "Authorization: Bearer {token}",
+        path = "/api/2.0/sql/warehouses")]
+    pub(crate) async fn warehouses(#[param] token: &str) -> feignhttp::Result<WarehouseResponse> {}
+
+    #[post(url = crate::BASE_URL,
+        headers = "Authorization: Bearer {token}",
+        path = "/api/2.0/sql/warehouses/{id}/stop")]
+    pub(crate) async fn stop_warehouse(#[param] id: &str, #[param] token: &str) -> feignhttp::Result<String> {}
+
+    #[post(url = crate::BASE_URL,
+        headers = "Authorization: Bearer {token}",
+        path = "/api/2.0/sql/warehouses/{id}/start")]
+    pub(crate) async fn start_warehouse(#[param] id: &str, #[param] token: &str) -> feignhttp::Result<String> {}
+
+    /// Rust representation of a list warehouses response
+    #[allow(unused)]
+    #[derive(Debug, Deserialize)]
+    pub struct WarehouseResponse {
+        pub warehouses: Vec<crate::sql::Warehouse>,
+    }
+}
+
+#[cfg(test)]
+mod tests {}