Update Athena example (#1007)

Signed-off-by: Kevin Su <[email protected]> Signed-off-by: Kevin Su <[email protected]> Co-authored-by: Iaroslav Ciupin <[email protected]>
flyteorg · Jul 6, 2023 · 73840e6 · 73840e6
1 parent a49ba51
commit 73840e6
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/examples/athena_plugin/athena_plugin/athena.py b/examples/athena_plugin/athena_plugin/athena.py
@@ -81,9 +81,9 @@ def no_io_wf():
     name="sql.athena.w_io",
     # Define inputs as well as their types that can be used to customize the query.
     inputs=kwtypes(iso_code=str),
-    task_config=AthenaConfig(database="vaccinations"),
+    task_config=AthenaConfig(workgroup="primary", catalog="AwsDataCatalog", database="vaccinations"),
     query_template="""
-    SELECT * FROM vaccinations where iso_code like  {{ .inputs.iso_code }}
+    SELECT * FROM vaccinations where iso_code like  '{{ .Inputs.iso_code }}'
     """,
     # While we define a generic schema as the output here, Flyte supports more strongly typed schemas to provide
     # better compile-time checks for task compatibility. Refer to :py:class:`flytekit.FlyteSchema` for more details
@@ -103,7 +103,7 @@ def no_io_wf():
 @task
 def manipulate_athena_schema(s: FlyteSchema) -> FlyteSchema:
     df = s.open().all()
-    return df[df.total_vaccinations.notnull()]
+    return df
 
 
 @workflow