feat: add s3 snippets (#2052)

* feat: add s3 snippets

* fix: rename to push pull aggregate

---------

Co-authored-by: Ruben Fiszel <ruben@windmill.dev>

backend/windmill-common/src/variables.rs

@@ -73,7 +73,7 @@ pub fn get_reserved_variables(
     flow_path: Option<String>,
     schedule_path: Option<String>,
     step_id: Option<String>,
-) -> [ContextualVariable; 14] {
+) -> [ContextualVariable; 15] {
     let state_path = {
         let flow_path = flow_path
             .clone()
@@ -117,6 +117,27 @@ pub fn get_reserved_variables(
         }
     };
 
+    let joined_schedule_path = schedule_path
+        .clone()
+        .unwrap_or("manual".to_string())
+        .split("/")
+        .collect::<Vec<&str>>()
+        .join("_");
+    let ts = chrono::Utc::now().timestamp_millis();
+    let object_path = if let Some(flow_path) = flow_path.clone() {
+        let flow_path = flow_path.split("/").collect::<Vec<&str>>().join("_");
+        let step_id = step_id.clone().unwrap_or("".to_string());
+        format!("{flow_path}/{joined_schedule_path}/{step_id}/{ts}_{job_id}")
+    } else {
+        let joined_script_path = path
+            .clone()
+            .unwrap_or("".to_string())
+            .split("/")
+            .collect::<Vec<&str>>()
+            .join("_");
+        format!("{joined_script_path}/{joined_schedule_path}/{ts}_{job_id}")
+    };
+
     [
         ContextualVariable {
             name: "WM_WORKSPACE".to_string(),
@@ -192,5 +213,10 @@ pub fn get_reserved_variables(
             value: step_id.unwrap_or_else(|| "".to_string()),
             description: "The node id in a flow (like 'a', 'b', or 'f')".to_string(),
         },
+        ContextualVariable {
+            name: "WM_OBJECT_PATH".to_string(),
+            value: object_path,
+            description: "Script or flow step execution unique path, useful for storing results in an external service".to_string(),
+        }
     ]
 }

frontend/src/lib/components/flows/content/FlowModuleComponent.svelte

@@ -7,7 +7,7 @@
 	import ModulePreview from '$lib/components/ModulePreview.svelte'
 	import { createScriptFromInlineScript, fork } from '$lib/components/flows/flowStateUtils'
 
-	import { RawScript, type FlowModule } from '$lib/gen'
+	import { RawScript, type FlowModule, Script } from '$lib/gen'
 	import FlowCard from '../common/FlowCard.svelte'
 	import FlowModuleHeader from './FlowModuleHeader.svelte'
 	import { getLatestHashForScript, scriptLangToEditorLang } from '$lib/scripts'
@@ -34,6 +34,10 @@
 	import { SecondsInput } from '$lib/components/common'
 	import DiffEditor from '$lib/components/DiffEditor.svelte'
 	import FlowModuleTimeout from './FlowModuleTimeout.svelte'
+	import HighlightCode from '$lib/components/HighlightCode.svelte'
+	import ToggleButtonGroup from '$lib/components/common/toggleButton-v2/ToggleButtonGroup.svelte'
+	import ToggleButton from '$lib/components/common/toggleButton-v2/ToggleButton.svelte'
+	import s3Scripts from './s3Scripts/lib'
 
 	const { selectedId, previewArgs, flowStateStore, flowStore, saveDraft } =
 		getContext<FlowEditorContext>('FlowEditorContext')
@@ -60,6 +64,7 @@
 	}
 	let selected = 'inputs'
 	let advancedSelected = 'retries'
+	let s3Kind = 'push'
 	let wrapper: HTMLDivElement
 	let panes: HTMLElement
 	let totalTopGap = 0
@@ -296,6 +301,9 @@
 										<Tab value="mock">Mock</Tab>
 										<Tab value="same_worker">Shared Directory</Tab>
 										<Tab value="timeout">Timeout</Tab>
+										{#if flowModule.value['language'] === 'python3' || flowModule.value['language'] === 'deno'}
+											<Tab value="s3">S3</Tab>
+										{/if}
 									{/if}
 								</Tabs>
 								<div class="h-[calc(100%-32px)] overflow-auto p-4">
@@ -368,6 +376,35 @@
 												}}>Set shared directory in the flow settings</Button
 											>
 										</div>
+									{:else if advancedSelected === 's3'}
+										<div>
+											<h2 class="pb-4">
+												S3 snippets
+												<Tooltip
+													>Pull, push and aggregate snippets for S3, particularly useful for ETL
+													processes.</Tooltip
+												>
+											</h2>
+										</div>
+										<div class="flex gap-2 justify-between mb-4 items-center">
+											<div class="flex gap-2">
+												<ToggleButtonGroup bind:selected={s3Kind} class="w-auto">
+													<ToggleButton value="push" size="sm" label="Push" />
+													<ToggleButton value="pull" size="sm" label="Pull" />
+													<ToggleButton value="aggregate" size="sm" label="Aggregate" />
+												</ToggleButtonGroup>
+											</div>
+											<Button
+												size="xs"
+												on:click={() =>
+													editor.setCode(s3Scripts[flowModule.value['language']][s3Kind])}
+												>Apply snippet
+											</Button>
+										</div>
+										<HighlightCode
+											language={Script.language.DENO}
+											code={s3Scripts[flowModule.value['language']][s3Kind]}
+										/>
 									{/if}
 								</div>
 							{/if}

frontend/src/lib/components/flows/content/s3Scripts/deno.ts

@@ -0,0 +1,105 @@
+const deno = {
+	push: `import { S3Client } from "https://deno.land/x/s3_lite_client@0.2.0/mod.ts";
+
+type S3 = {
+  port: number;
+  bucket: string;
+  region: string;
+  useSSL: boolean;
+  endPoint: string;
+  accessKey: string;
+  pathStyle: boolean;
+  secretKey: string;
+};
+
+export async function main(
+  s3Config: S3,
+  basePath = "windmill",
+  objectName: string,
+  data: string | Uint8Array | ReadableStream<Uint8Array>,
+) {
+  // flow_path/schedule_path_or_manual/flow_step_id/ts_job_id
+  const objectPath = Deno.env.get("WM_OBJECT_PATH"); 
+
+  const fullPath = basePath + "/" + objectPath + "/" + objectName;
+
+  const s3Client = new S3Client(s3Config);
+
+  await s3Client.putObject(fullPath, data);
+
+  return fullPath;
+}`,
+	pull: `import { S3Client } from "https://deno.land/x/s3_lite_client@0.2.0/mod.ts";
+
+type S3 = {
+  port: number;
+  bucket: string;
+  region: string;
+  useSSL: boolean;
+  endPoint: string;
+  accessKey: string;
+  pathStyle: boolean;
+  secretKey: string;
+};
+
+export async function main(
+  s3Config: S3,
+  objectPath: string,
+) {
+
+  const s3Client = new S3Client(s3Config);
+
+  const response = await s3Client.getObject(objectPath)
+  // for instance, if it is a text file
+  const result = await response.text()
+  return result
+}`,
+	aggregate: `import { S3Client } from "https://deno.land/x/s3_lite_client@0.2.0/mod.ts";
+
+type S3 = {
+  port: number;
+  bucket: string;
+  region: string;
+  useSSL: boolean;
+  endPoint: string;
+  accessKey: string;
+  pathStyle: boolean;
+  secretKey: string;
+};
+
+export async function main(
+  s3Config: S3,
+  objectPath: string,
+  last_n = 10,
+) {
+
+  // object path assumed to be of the form windmill/flow_path/schedule_path_or_manual/flow_step_id/ts_job_id/**
+  const prefix = objectPath.split("/").slice(0, 4).join("/")
+
+  const s3Client = new S3Client(s3Config);
+  
+  // will return the object keys of the last_n jobs
+  const objs = {};
+  for await (const entry of s3Client.listObjects({ prefix })) {
+    const obj_key = entry.key
+    const ts = parseInt(obj_key.split("/")[4].split("_")[0])
+    if (ts in objs) {
+      objs[ts].append()
+    } else {
+      objs[ts] = [obj_key]
+    }
+  }
+
+  const tss = Object.keys(objs).sort().slice(-last_n)
+
+
+  const final_objs = []
+  for (const ts of tss) {
+    final_objs.push(...objs[ts])
+  }
+
+  return final_objs;
+}`
+}
+
+export default deno

frontend/src/lib/components/flows/content/s3Scripts/lib.ts

@@ -0,0 +1,9 @@
+import deno from './deno'
+import python3 from './python3'
+
+const s3Scripts = {
+	deno,
+	python3
+}
+
+export default s3Scripts

frontend/src/lib/components/flows/content/s3Scripts/python3.ts

@@ -0,0 +1,118 @@
+const python3 = {
+	push: `import os
+import boto3
+from typing import TypedDict, Union
+
+class s3(TypedDict):
+    port: int
+    bucket: str
+    region: str
+    useSSL: bool
+    endPoint: str
+    accessKey: str
+    pathStyle: bool
+    secretKey: str
+
+def main(
+    s3_config: s3,
+    object_name: str,
+    data: Union[str, bytes],
+    base_path: str = "windmill",
+):
+    # flow_path/schedule_path_or_manual/flow_step_id/ts_job_id
+    object_path = os.getenv("WM_OBJECT_PATH")
+
+    full_path = base_path + "/" + object_path + "/" + object_name
+
+    s3Client = boto3.client(
+        's3',
+        region_name=s3_config['region'],
+        aws_access_key_id=s3_config['accessKey'],
+        aws_secret_access_key=s3_config['secretKey']
+    )
+
+    s3Client.put_object(Body=data, Bucket=s3_config['bucket'], Key=full_path)
+
+    return full_path`,
+	pull: `import boto3
+from typing import TypedDict
+
+class s3(TypedDict):
+    port: int
+    bucket: str
+    region: str
+    useSSL: bool
+    endPoint: str
+    accessKey: str
+    pathStyle: bool
+    secretKey: str
+
+def main(
+    s3_config: s3,
+    object_path: str,
+):
+
+    s3Client = boto3.client(
+        's3',
+        region_name=s3_config['region'],
+        aws_access_key_id=s3_config['accessKey'],
+        aws_secret_access_key=s3_config['secretKey']
+    )
+
+    obj = s3Client.get_object(Bucket=s3_config["bucket"], Key=object_path)["Body"].read()
+
+    # for instance, if it is a text file
+    return str(obj, encoding="utf-8")`,
+	aggregate: `import os
+import boto3
+
+from typing import TypedDict
+
+class s3(TypedDict):
+    port: int
+    bucket: str
+    region: str
+    useSSL: bool
+    endPoint: str
+    accessKey: str
+    pathStyle: bool
+    secretKey: str
+
+
+def main(
+    s3_config: s3,
+    object_path: str,
+    last_n = 10,
+):
+    # object path assumed to be of the form windmill/flow_path/schedule_path_or_manual/flow_step_id/ts_job_id/**
+    prefix = "/".join(object_path.split("/")[:4])
+
+    s3Client = boto3.client(
+        "s3",
+        region_name=s3_config["region"],
+        aws_access_key_id=s3_config["accessKey"],
+        aws_secret_access_key=s3_config["secretKey"],
+    )
+
+    # will return the object keys of the last_n jobs
+    objs = {}
+    for content in s3Client.list_objects(
+        Bucket=s3_config["bucket"], Prefix=prefix
+    )["Contents"]:
+        obj_key = content["Key"]
+        ts = int(obj_key.split("/")[4].split("_")[0])
+        if ts in objs:
+            objs[ts].append(obj_key)
+        else:
+            objs[ts] = [obj_key]
+
+    tss = sorted(objs.keys(), reverse=True)[:last_n]
+
+    final_objs = []
+    for ts in tss:
+        final_objs.extend(objs[ts])
+
+    return final_objs`
+}
+
+export default python3