From d63f5ee26e41ee763505b5fda220c5a30865c60e Mon Sep 17 00:00:00 2001
From: Ruben Fiszel <ruben@rubenfiszel.com>
Date: Sat, 23 Mar 2024 18:42:38 +0100
Subject: [PATCH] fix: add support for azure storage on distributed cache/logs

---
 backend/windmill-api/openapi.yaml             |  31 +++-
 backend/windmill-api/src/settings.rs          |   9 +-
 .../windmill-worker/src/python_executor.rs    |   2 +-
 .../ObjectStoreConfigSettings.svelte          | 149 ++++++++++++++----
 .../src/lib/components/TestConnection.svelte  |   2 +-
 .../src/lib/components/instanceSettings.ts    |   2 +-
 6 files changed, 149 insertions(+), 46 deletions(-)

diff --git a/backend/windmill-api/openapi.yaml b/backend/windmill-api/openapi.yaml
index 34b9027bd2..130e1014e9 100644
--- a/backend/windmill-api/openapi.yaml
+++ b/backend/windmill-api/openapi.yaml
@@ -718,24 +718,47 @@ paths:
               schema:
                 type: string
 
-  /settings/test_s3_config:
+
+    # pub use_ssl: Option<bool>,
+    # #[serde(rename = "accountName")]
+    # pub account_name: String,
+    # #[serde(rename = "tenantId")]
+    # pub tenant_id: Option<String>,
+    # #[serde(rename = "clientId")]
+    # pub client_id: Option<String>,
+    # #[serde(rename = "containerName")]
+    # pub container_name: String,
+    # #[serde(rename = "accessKey")]
+    # pub access_key: Option<String>,
+
+  /settings/test_object_storage_config:
     post:
-      summary: test s3 config
-      operationId: testS3Config
+      summary: test object storage config
+      operationId: testObjectStorageConfig
       tags:
         - setting
       requestBody:
-        description: test s3 config
+        description: test object storage config
         required: true
         content:
           application/json:
             schema:
               type: object
               properties:
+                type:
+                  type: string
                 bucket:
                   type: string
                 region:
                   type: string
+                accessKey:
+                  type: string
+                accountName:
+                  type: string
+                tenantId:
+                  type: string
+                clientId:
+                  type: string
                 access_key:
                   type: string
                 secret_key:
diff --git a/backend/windmill-api/src/settings.rs b/backend/windmill-api/src/settings.rs
index 43c144066b..197b529014 100644
--- a/backend/windmill-api/src/settings.rs
+++ b/backend/windmill-api/src/settings.rs
@@ -44,7 +44,7 @@ pub fn global_service() -> Router {
 
     #[cfg(feature = "parquet")]
     {
-        return r.route("/test_s3_config", post(test_s3_bucket));
+        return r.route("/test_object_storage_config", post(test_s3_bucket));
     }
 
     #[cfg(not(feature = "parquet"))]
@@ -126,7 +126,7 @@ pub async fn test_s3_bucket(
         "/test-s3-bucket-{uuid}",
         uuid = uuid::Uuid::new_v4()
     ));
-    tracing::info!("Testing s3 bucket at path: {path}");
+    tracing::info!("Testing blob storage at path: {path}");
     client
         .put(&path, Bytes::from_static(b"hello"))
         .await
@@ -140,13 +140,14 @@ pub async fn test_s3_bucket(
         .map_err(to_anyhow)?;
     if content != Bytes::from_static(b"hello") {
         return Err(error::Error::InternalErr(
-            "Failed to read back from s3".to_string(),
+            "Failed to read back from blob storage".to_string(),
         ));
     }
     client.delete(&path).await.map_err(to_anyhow)?;
-    Ok("Tested bucket successfully".to_string())
+    Ok("Tested blob storage successfully".to_string())
 }
 
+
 #[derive(Deserialize)]
 pub struct TestKey {
     pub license_key: String,
diff --git a/backend/windmill-worker/src/python_executor.rs b/backend/windmill-worker/src/python_executor.rs
index eab0d419f2..fbcae34469 100644
--- a/backend/windmill-worker/src/python_executor.rs
+++ b/backend/windmill-worker/src/python_executor.rs
@@ -860,7 +860,7 @@ pub async fn handle_python_reqs(
                     job_id.clone(),
                     w_id.to_string(),
                     format!(
-                        "pulled {} from s3 cache in {}ms",
+                        "pulled {} from distributed cache in {}ms",
                         pulled.join(", "),
                         start.elapsed().as_millis()
                     ),
diff --git a/frontend/src/lib/components/ObjectStoreConfigSettings.svelte b/frontend/src/lib/components/ObjectStoreConfigSettings.svelte
index d0814ea8f8..fa1734abc2 100644
--- a/frontend/src/lib/components/ObjectStoreConfigSettings.svelte
+++ b/frontend/src/lib/components/ObjectStoreConfigSettings.svelte
@@ -1,12 +1,12 @@
 <script lang="ts">
 	import { Database, Loader2 } from 'lucide-svelte'
 	import Toggle from './Toggle.svelte'
-	import { Button } from './common'
+	import { Button, Tab, Tabs } from './common'
 	import { SettingService } from '$lib/gen'
 	import { sendUserToast } from '$lib/toast'
 	import TestConnection from './TestConnection.svelte'
 
-	type BucketConfig = {
+	type S3Config = {
 		type: 'S3'
 		bucket: string
 		region: string
@@ -14,7 +14,19 @@
 		secret_key: string
 		endpoint: string
 	}
-	export let bucket_config: BucketConfig | undefined = undefined
+
+	type AzureConfig = {
+		type: 'Azure'
+		accountName: string
+		containerName: string
+		useSSL?: boolean
+		tenantId: string
+		clientId: string
+		accessKey: string
+		endpoint?: string
+	}
+
+	export let bucket_config: S3Config | AzureConfig | undefined = undefined
 
 	let loading = false
 
@@ -22,7 +34,7 @@
 		loading = true
 		try {
 			if (bucket_config) {
-				await SettingService.testS3Config({ requestBody: bucket_config })
+				await SettingService.testObjectStorageConfig({ requestBody: bucket_config })
 				sendUserToast('Connection successful', false)
 			}
 		} catch (e) {
@@ -77,36 +89,103 @@
 			buttonTextOverride="Test from a worker"
 		/>
 	</div>
-	<label class="block pb-2">
-		<span class="text-primary font-semibold text-sm">Bucket</span>
-		<input type="text" placeholder="bucket-name" bind:value={bucket_config.bucket} />
-	</label>
+	<Tabs
+		bind:selected={bucket_config.type}
+		on:selected={(e) => {
+			if (e.detail === 'S3') {
+				bucket_config = {
+					type: 'S3',
+					bucket: '',
+					region: '',
+					access_key: '',
+					secret_key: '',
+					endpoint: ''
+				}
+			} else if (e.detail === 'Azure') {
+				bucket_config = {
+					type: 'Azure',
+					accountName: '',
+					containerName: '',
+					useSSL: false,
+					tenantId: '',
+					clientId: '',
+					accessKey: ''
+				}
+			}
+		}}
+	>
+		<Tab size="sm" value="S3">S3</Tab>
+		<Tab size="sm" value="Azure">Azure Blob</Tab>
+	</Tabs>
+	{#if bucket_config.type === 'S3'}
+		<label class="block pb-2">
+			<span class="text-primary font-semibold text-sm">Bucket</span>
+			<input type="text" placeholder="bucket-name" bind:value={bucket_config.bucket} />
+		</label>
 
-	<label class="block pb-2">
-		<span class="text-primary font-semibold text-sm">Region</span>
-		<span class="text-tertiary text-2xs"
-			>If left empty, will be derived automatically from $AWS_REGION</span
-		>
-		<input type="text" bind:value={bucket_config.region} />
-	</label>
-	<label class="block pb-2">
-		<span class="text-primary font-semibold text-sm">Access Key ID</span>
-		<span class="text-tertiary text-2xs"
-			>If left empty, will be derived automatically from $AWS_ACCESS_KEY_ID, pod or ec2 profile</span
-		>
-		<input type="text" bind:value={bucket_config.access_key} />
-	</label>
-	<label class="block pb-2">
-		<span class="text-primary font-semibold text-sm">Secret Key</span>
-		<span class="text-tertiary text-2xs"
-			>If left empty, will be derived automatically from $AWS_SECRET_KEY, pod or ec2 profile</span
-		>
-		<input type="text" bind:value={bucket_config.secret_key} />
-	</label>
-	<label class="block pb-2">
-		<span class="text-primary font-semibold text-sm">Endpoint</span>
-		<span class="text-tertiary text-2xs">Only needed for non AWS S3 providers like R2 or MinIo</span
-		>
-		<input type="text" bind:value={bucket_config.endpoint} />
-	</label>
+		<label class="block pb-2">
+			<span class="text-primary font-semibold text-sm">Region</span>
+			<span class="text-tertiary text-2xs"
+				>If left empty, will be derived automatically from $AWS_REGION</span
+			>
+			<input type="text" bind:value={bucket_config.region} />
+		</label>
+		<label class="block pb-2">
+			<span class="text-primary font-semibold text-sm">Access Key ID</span>
+			<span class="text-tertiary text-2xs"
+				>If left empty, will be derived automatically from $AWS_ACCESS_KEY_ID, pod or ec2 profile</span
+			>
+			<input type="text" bind:value={bucket_config.access_key} />
+		</label>
+		<label class="block pb-2">
+			<span class="text-primary font-semibold text-sm">Secret Key</span>
+			<span class="text-tertiary text-2xs"
+				>If left empty, will be derived automatically from $AWS_SECRET_KEY, pod or ec2 profile</span
+			>
+			<input type="password" bind:value={bucket_config.secret_key} />
+		</label>
+		<label class="block pb-2">
+			<span class="text-primary font-semibold text-sm">Endpoint</span>
+			<span class="text-tertiary text-2xs"
+				>Only needed for non AWS S3 providers like R2 or MinIo</span
+			>
+			<input type="text" bind:value={bucket_config.endpoint} />
+		</label>
+	{:else if bucket_config.type === 'Azure'}
+		<label class="block pb-2">
+			<span class="text-primary font-semibold text-sm">Account Name</span>
+			<input type="text" placeholder="account-name" bind:value={bucket_config.accountName} />
+		</label>
+		<label class="block pb-2">
+			<span class="text-primary font-semibold text-sm">Container Name</span>
+			<input type="text" placeholder="container-name" bind:value={bucket_config.containerName} />
+		</label>
+		<label class="block pb-2">
+			<span class="text-primary font-semibold text-sm">Access Key</span>
+			<input type="password" bind:value={bucket_config.accessKey} />
+		</label>
+		<label class="block pb-2">
+			<span class="text-primary font-semibold text-sm"
+				>Tenant ID <span class="text-2xs text-tertiary">(optional)</span></span
+			>
+			<input type="text" bind:value={bucket_config.tenantId} />
+		</label>
+		<label class="block pb-2">
+			<span class="text-primary font-semibold text-sm"
+				>Client ID <span class="text-2xs text-tertiary">(optional)</span></span
+			>
+			<input type="text" bind:value={bucket_config.clientId} />
+		</label>
+		<label class="block pb-2">
+			<span class="text-primary font-semibold text-sm"
+				>Endpoint <span class="text-2xs text-tertiary">(optional)</span></span
+			>
+			<span class="text-tertiary text-2xs"
+				>Only needed for non Azure Blob providers like Azurite</span
+			>
+			<input type="text" bind:value={bucket_config.endpoint} />
+		</label>
+	{:else}
+		<div>Unknown bucket type {bucket_config['type']}</div>
+	{/if}
 {/if}
diff --git a/frontend/src/lib/components/TestConnection.svelte b/frontend/src/lib/components/TestConnection.svelte
index 633e9eeb6f..cda3ccfdb9 100644
--- a/frontend/src/lib/components/TestConnection.svelte
+++ b/frontend/src/lib/components/TestConnection.svelte
@@ -90,7 +90,7 @@ export async function main(s3: S3) {
 const process = require('process');
 
 export async function main(bucket: any) {
-	const req = await fetch(process.env.BASE_URL + '/api/settings/test_s3_config', {
+	const req = await fetch(process.env.BASE_URL + '/api/settings/test_object_storage_config', {
 		method: 'POST',
 		headers: {
 			'Content-Type': 'application/json',
diff --git a/frontend/src/lib/components/instanceSettings.ts b/frontend/src/lib/components/instanceSettings.ts
index 0dfb868121..d083c0a271 100644
--- a/frontend/src/lib/components/instanceSettings.ts
+++ b/frontend/src/lib/components/instanceSettings.ts
@@ -95,7 +95,7 @@ export const settings: Record<string, Setting[]> = {
 			ee_only: ''
 		},
 		{
-			label: 'S3 for Python Cache & Large Logs',
+			label: 'S3/Azure for Python Cache & Large Logs',
 			description: 'Bucket to store large logs and cache for distributed python jobs.',
 			key: 'object_store_cache_config',
 			fieldType: 'object_store_config',