From 8d87712e2cf6753dc99108434150676ee0afa313 Mon Sep 17 00:00:00 2001
From: Guillaume Bouvignies <guillaume@windmill.dev>
Date: Thu, 14 Dec 2023 15:35:24 +0100
Subject: [PATCH] chore: Update S3 code snippets (#2854)

* chore: Update S3 code snippets

* update tooltip

* cleanup code
---
 .../flows/content/FlowModuleComponent.svelte  |  17 +-
 .../flows/content/s3Scripts/deno.ts           | 116 ++-------
 .../flows/content/s3Scripts/python3.ts        | 226 ++++++++++--------
 3 files changed, 156 insertions(+), 203 deletions(-)

diff --git a/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte b/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
index 93102e5d10..f1d2b3e6c7 100644
--- a/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
+++ b/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
@@ -72,7 +72,7 @@
 	let selected = 'inputs'
 	let advancedSelected = 'retries'
 	let advancedRuntimeSelected = 'concurrency'
-	let s3Kind = 'push'
+	let s3Kind = 's3_client'
 	let wrapper: HTMLDivElement
 	let panes: HTMLElement
 	let totalTopGap = 0
@@ -498,19 +498,24 @@
 											<h2 class="pb-4">
 												S3 snippets
 												<Tooltip>
-													Pull, push and aggregate snippets for S3, particularly useful for ETL
-													processes.
+													Read/Write object from/to S3 and leverage Polars and DuckDB to run
+													efficient 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" />
+													{#if flowModule.value['language'] === 'deno'}
+														<ToggleButton value="s3_client" size="sm" label="S3 lite client" />
+													{:else}
+														<ToggleButton value="s3_client" size="sm" label="Boto3" />
+														<ToggleButton value="polars" size="sm" label="Polars" />
+														<ToggleButton value="duckdb" size="sm" label="DuckDB" />
+													{/if}
 												</ToggleButtonGroup>
 											</div>
+
 											<Button
 												size="xs"
 												on:click={() =>
diff --git a/frontend/src/lib/components/flows/content/s3Scripts/deno.ts b/frontend/src/lib/components/flows/content/s3Scripts/deno.ts
index 905b5ddb98..67e84f3a2a 100644
--- a/frontend/src/lib/components/flows/content/s3Scripts/deno.ts
+++ b/frontend/src/lib/components/flows/content/s3Scripts/deno.ts
@@ -1,105 +1,35 @@
 const deno = {
-	push: `import { S3Client } from "https://deno.land/x/s3_lite_client@0.2.0/mod.ts";
+  s3_client: `import * as wmill from "npm:windmill-client@1";
+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;
-};
+type s3object = object;
 
-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"); 
+export async function main(inputFile: s3object) {
+  const s3Resource = await wmill.getResource(
+    "<PATH_TO_S3_RESOURCE>",
+  );
+  const s3Client = new S3Client(s3Resource);
 
-  const fullPath = basePath + "/" + objectPath + "/" + objectName;
+  const outputFile = "output/hello.txt"
 
-  const s3Client = new S3Client(s3Config);
+  // read object from S3
+  const getObjectResponse = await s3Client.getObject(inputFile["s3"]);
+  const inputObjContent = await getObjectResponse.text();
+  console.log(inputObjContent);
 
-  await s3Client.putObject(fullPath, data);
+  // write object to S3
+  await s3Client.putObject(outputFile, "Hello Windmill!");
 
-  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]
-    }
+  // list objects from bucket
+  for await (const obj of s3Client.listObjects({ prefix: "output/" })) {
+    console.log(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;
-}`
+  return {
+    "s3": outputFile,
+  };
+}
+`
 }
 
 export default deno
diff --git a/frontend/src/lib/components/flows/content/s3Scripts/python3.ts b/frontend/src/lib/components/flows/content/s3Scripts/python3.ts
index f2b7f42c06..c666939eb7 100644
--- a/frontend/src/lib/components/flows/content/s3Scripts/python3.ts
+++ b/frontend/src/lib/components/flows/content/s3Scripts/python3.ts
@@ -1,118 +1,136 @@
 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
+	s3_client: `import wmill
 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
+s3object = dict
 
 
-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(
+def main(input_file: s3object):
+    s3_resource = wmill.get_resource("<PATH_TO_S3_RESOURCE>")
+    bucket = s3_resource["bucket"]
+    s3client = boto3.client(
         "s3",
-        region_name=s3_config["region"],
-        aws_access_key_id=s3_config["accessKey"],
-        aws_secret_access_key=s3_config["secretKey"],
+        region_name=s3_resource["region"],
+        aws_access_key_id=s3_resource["accessKey"],
+        aws_secret_access_key=s3_resource["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]
+    output_file = "output/hello.txt"
 
-    tss = sorted(objs.keys(), reverse=True)[:last_n]
+    # read object from S3 and print its content
+    input_obj = s3client.get_object(Bucket=bucket, Key=input_file["s3"])["Body"].read()
+    print(input_obj)
 
-    final_objs = []
-    for ts in tss:
-        final_objs.extend(objs[ts])
+    # write object to s3
+    s3client.put_object(Bucket=bucket, Key=output_file, Body="Hello Windmill!")
 
-    return final_objs`
+    # download file to the job temporary folder:
+    s3client.download_file(
+        Bucket=bucket, Key=input_file["s3"], Filename="./download.txt"
+    )
+    with open("./download.txt", mode="rb") as downloaded_file:
+        print(downloaded_file.read())
+
+    # upload file from temporary folder to S3
+    uploaded_file = "output/uploaded.txt"
+    with open("./upload.txt", mode="wb") as file_to_upload:
+        file_to_upload.write(str.encode("Hello Windmill!"))
+    s3client.upload_file(Bucket=bucket, Key=uploaded_file, Filename="./upload.txt")
+
+    # see https://boto3.amazonaws.com/v1/documentation/api/latest/guide/s3-examples.html
+    # and https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3.html
+    # for more code examples (listing object, deleting files, etc)
+    return [
+        s3object({"s3": output_file}),
+        s3object({"s3": uploaded_file}),
+    ]
+`,
+
+    polars: `import wmill
+import polars as pl
+import s3fs
+
+s3object = dict
+
+
+def main(input_file: s3object):
+    s3 = s3fs.S3FileSystem(
+        # this will default to the workspace s3 resource
+        **wmill.polars_connection_settings()["s3fs_args"]
+        # this will use the designated resource
+        # **wmill.polars_connection_settings("<PATH_TO_S3_RESOURCE>")["s3fs_args"]
+    )
+
+    bucket = "<S3_BUCKET_NAME>"
+    input_uri = "s3://{}/{}".format(bucket, input_file["s3"])
+    output_file = "output/result.parquet"
+    output_uri = "s3://{}/{}".format(bucket, output_file)
+
+    with (
+        s3.open(input_uri, mode="rb") as input_s3,
+        s3.open(output_uri, mode="wb") as output_s3,
+    ):
+        # input is a parquet file, we use read_parquet in lazy mode.
+        # Polars can read various file types, see
+        # https://pola-rs.github.io/polars/py-polars/html/reference/io.html
+        input_df = pl.read_parquet(input_s3).lazy()
+
+        # process the Polars dataframe. See Polars docs:
+        # for dataframe: https://pola-rs.github.io/polars/py-polars/html/reference/dataframe/index.html
+        # for lazy dataframe: https://pola-rs.github.io/polars/py-polars/html/reference/lazyframe/index.html
+        output_df = input_df.collect()
+        print(output_df)
+
+        # persist the output dataframe back to S3 and return it
+        output_df.write_parquet(output_s3)
+
+    return s3object({"s3": output_file})
+`,
+
+    duckdb: `import wmill
+import duckdb
+
+s3object = dict
+
+
+def main(input_file: s3object):
+    # create a DuckDB database in memory
+    # see https://duckdb.org/docs/api/python/dbapi
+    conn = duckdb.connect()
+    # connect duck db to the S3 bucket - this will default to the workspace s3 resource
+    conn.execute(wmill.duckdb_connection_settings()["connection_settings_str"])
+    # this will use the designated resource
+    # conn.execute(wmill.duckdb_connection_settings("<PATH_TO_S3_RESOURCE>")["connection_settings_str"])
+
+    bucket = "<S3_BUCKET_NAME>"
+    input_uri = "s3://{}/{}".format(bucket, input_file["s3"])
+    output_file = "output/result.parquet"
+    output_uri = "s3://{}/{}".format(bucket, output_file)
+
+    # Run queries directly on the parquet file
+    query_result = conn.sql(
+        """
+        SELECT * FROM read_parquet('{}')
+    """.format(
+            input_uri
+        )
+    )
+    query_result.show()
+
+    # Write the result of a query to a different parquet file on S3
+    conn.execute(
+        """
+        COPY (
+            SELECT COUNT(*) FROM read_parquet('{input_uri}')
+        ) TO '{output_uri}' (FORMAT 'parquet');
+    """.format(
+            input_uri=input_uri, output_uri=output_uri
+        )
+    )
+
+    conn.close()
+    return s3object({"s3": output_file})
+`,
 }
 
 export default python3