[PLUGINS] Bump Version [snowflake]

plugins/snowflake/superduper_snowflake/__init__.py

@@ -1,7 +1,7 @@
 from .vector_search import SnowflakeVectorSearcher as VectorSearcher
 from .data_backend import SnowflakeDataBackend as DataBackend
 
-__version__ = "0.5.12"
+__version__ = "0.5.13"
 
 __all__ = [
     "VectorSearcher",

plugins/snowflake/superduper_snowflake/data_backend.py

@@ -78,6 +78,10 @@ def _connection_callback(self, uri):
             return IbisDataBackend._connection_callback(uri)
         return ibis.snowflake.from_connection(self._do_connection_callback(uri), create_object_udfs=False), 'snowflake', False
 
+    def reconnect(self):
+        super().reconnect()
+        self.snowpark = self._get_snowpark_session(self.uri)
+
     def insert(self, table_name, raw_documents):
         ibis_schema = self.conn.table(table_name).schema()
         df = pandas.DataFrame(raw_documents)

plugins/snowflake/superduper_snowflake/vector_search.py

@@ -1,3 +1,4 @@
+from functools import wraps
 import os
 import re
 import typing as t
@@ -10,6 +11,18 @@
     from superduper.components.vector_index import VectorIndex
 
 
+def retry(f):
+    @wraps(f)
+    def wrapper(self, *args, **kwargs):
+        try:
+            return f(self, *args, **kwargs)
+        except Exception as e:
+            if 'token' in str(e):
+                self.session = SnowflakeVectorSearcher.create_session(CFG.data_backend)
+            return f(self, *args, **kwargs)
+    return wrapper
+
+
 class SnowflakeVectorSearcher(BaseVectorSearcher):
     """Vector searcher implementation of atlas vector search.
 
@@ -161,6 +174,7 @@ def find_nearest_from_id(self, id: str, n=100, within_ids=None):
         ).collect()
         return self.find_nearest_from_array(result, n=n, within_ids=within_ids)
 
+    @retry
     def find_nearest_from_array(self, h, n=100, within_ids=None):
         """Find the nearest vectors to the given vector.