From 4211c33416c27559b9ae2ded5b1d3bb310691018 Mon Sep 17 00:00:00 2001
From: Matt Kornfield <kornfield@gretel.ai>
Date: Tue, 8 Aug 2023 08:45:05 -0700
Subject: [PATCH] Use get_artifact_handle instead of smart_open directly

---
 requirements.txt                            | 2 +-
 src/gretel_trainer/relational/sdk_extras.py | 5 ++---
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 2062012f..0ee5a4fc 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,6 @@
 boto3~=1.20
 dask[dataframe]==2023.5.1
-gretel-client>=0.16.0
+gretel-client>=0.16.12
 jinja2~=3.1
 networkx~=3.0
 numpy~=1.20
diff --git a/src/gretel_trainer/relational/sdk_extras.py b/src/gretel_trainer/relational/sdk_extras.py
index e9550bad..2c6f02f2 100644
--- a/src/gretel_trainer/relational/sdk_extras.py
+++ b/src/gretel_trainer/relational/sdk_extras.py
@@ -52,9 +52,8 @@ def download_file_artifact(
         artifact_name: str,
         out_path: Union[str, Path],
     ) -> None:
-        download_link = gretel_object.get_artifact_link(artifact_name)
         try:
-            with smart_open.open(download_link, "rb") as src, smart_open.open(
+            with gretel_object.get_artifact_handle() as src, smart_open.open(
                 out_path, "wb"
             ) as dest:
                 shutil.copyfileobj(src, dest)
@@ -80,7 +79,7 @@ def sqs_score_from_full_report(self, report: dict[str, Any]) -> Optional[int]:
                     return field_dict["value"]
 
     def get_record_handler_data(self, record_handler: RecordHandler) -> pd.DataFrame:
-        with smart_open.open(record_handler.get_artifact_link("data"), "rb") as data:
+        with record_handler.get_artifact_handle() as data:
             return pd.read_csv(data)
 
     def start_job_if_possible(