Use get_artifact_handle instead of smart_open directly (#153)

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

src/gretel_trainer/relational/sdk_extras.py

@@ -52,11 +52,10 @@ 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(
-                out_path, "wb"
-            ) as dest:
+            with gretel_object.get_artifact_handle(
+                artifact_name
+            ) as src, smart_open.open(out_path, "wb") as dest:
                 shutil.copyfileobj(src, dest)
         except:
             logger.warning(f"Failed to download `{artifact_name}`")
@@ -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("data") as data:
             return pd.read_csv(data)
 
     def start_job_if_possible(

tests/relational/test_multi_table_restore.py

@@ -148,7 +148,7 @@ def create_backup(
 
 def get_local_name(artifact_id):
     local_name = None
-    for key, pointers in ARTIFACTS.items():
+    for _, pointers in ARTIFACTS.items():
         if pointers["artifact_id"] == artifact_id:
             local_name = pointers["local_name"]
     if local_name is None:
@@ -163,6 +163,13 @@ def get_artifact_link(artifact_id):
     return get_artifact_link
 
 
+def make_mock_get_artifact_handle(setup_path: Path):
+    def get_artifact_handle(artifact_id):
+        return open(setup_path / get_local_name(artifact_id), "rb")
+
+    return get_artifact_handle
+
+
 def make_mock_download_tar_artifact(setup_path: Path, working_path: Path):
     def download_tar_artifact(project, name, out_path):
         local_name = get_local_name(name)
@@ -193,6 +200,7 @@ def make_mock_model(
     model.model_id = name
     model.data_source = ARTIFACTS[f"train_{name}"]["artifact_id"]
     model.get_artifact_link = make_mock_get_artifact_link(setup_path)
+    model.get_artifact_handle = make_mock_get_artifact_handle(setup_path)
     model.get_record_handler.return_value = record_handler
     return model
 
@@ -329,6 +337,7 @@ def configure_mocks(
     models: dict[str, Mock] = {},
 ) -> None:
     project.get_artifact_link = make_mock_get_artifact_link(setup_path)
+    project.get_artifact_handle = make_mock_get_artifact_handle(setup_path)
     project.get_model = make_mock_get_model(models)
     download_tar_artifact.side_effect = make_mock_download_tar_artifact(
         setup_path,