Close datasets properly

snowfakery/data_generator_runtime.py

@@ -3,6 +3,7 @@
 from contextlib import contextmanager
 
 from typing import Optional, Dict, Sequence, Mapping, NamedTuple, Set
+from warnings import warn
 
 import jinja2
 import yaml
@@ -354,7 +355,10 @@ def __enter__(self, *args):
 
     def __exit__(self, *args):
         for plugin in self.plugin_instances.values():
-            plugin.close()
+            try:
+                plugin.close()
+            except Exception as e:
+                warn(f"Could not close {plugin} because {e}")
 
 
 class RuntimeContext:

snowfakery/standard_plugins/Salesforce.py

@@ -377,8 +377,19 @@ def _load_dataset(self, iteration_mode, rootpath, kwargs):
         return self.iterator
 
     def close(self):
-        self.iterator.close()
-        self.tempdir.close()
+        if self.iterator:
+            self.iterator.close()
+            self.iterator = None
+
+        if self.tempdir:
+            self.tempdir.cleanup()
+            self.tempdir = None
+
+    def __del__(self):
+        # in case close was not called
+        # properly, try to do an orderly
+        # cleanup
+        self.close()
 
 
 def create_tempfile_sql_db_iterator(mode, fieldnames, results):

snowfakery/standard_plugins/datasets.py

@@ -224,6 +224,10 @@ def iterate(self, **kwargs):
         def shuffle(self, **kwargs):
             return self.context.plugin.dataset_impl._iterate(self, "shuffle", kwargs)
 
+    def close(self):
+        if self.dataset_impl:
+            self.dataset_impl.close()
+
 
 class Dataset(DatasetPluginBase):
     def __init__(self, *args, **kwargs):