fix: Fix SQLite import issue (#4294)

franciscojavierarceo · franciscojavierarceo · commit 8d9b60e38016 · 2024-07-03T22:42:31.000-04:00
adding try and except block for import

Co-authored-by: Francisco Javier Arceo &lt;franciscojavierarceo@users.noreply.github.com&gt;
diff --git a/sdk/python/feast/infra/online_stores/sqlite.py b/sdk/python/feast/infra/online_stores/sqlite.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import itertools
+import logging
 import os
 import sqlite3
 import struct
@@ -20,7 +21,6 @@
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple, Union
 
-import sqlite_vec
 from google.protobuf.internal.containers import RepeatedScalarFieldContainer
 from pydantic import StrictStr
 
@@ -84,7 +84,9 @@ def _get_conn(self, config: RepoConfig):
         if not self._conn:
             db_path = self._get_db_path(config)
             self._conn = _initialize_conn(db_path)
-            if sys.version_info[0:2] == (3, 10):
+            if sys.version_info[0:2] == (3, 10) and config.online_store.vec_enabled:
+                import sqlite_vec  # noqa: F401
+
                 self._conn.enable_load_extension(True)  # type: ignore
                 sqlite_vec.load(self._conn)
 
@@ -410,6 +412,10 @@ def retrieve_online_documents(
 
 
 def _initialize_conn(db_path: str):
+    try:
+        import sqlite_vec  # noqa: F401
+    except ModuleNotFoundError:
+        logging.warning("Cannot use sqlite_vec for vector search")
     Path(db_path).parent.mkdir(exist_ok=True)
     return sqlite3.connect(
         db_path,
@@ -482,8 +488,13 @@ def from_proto(sqlite_table_proto: SqliteTableProto) -> Any:
 
     def update(self):
         if sys.version_info[0:2] == (3, 10):
-            self.conn.enable_load_extension(True)
-            sqlite_vec.load(self.conn)
+            try:
+                import sqlite_vec  # noqa: F401
+
+                self.conn.enable_load_extension(True)
+                sqlite_vec.load(self.conn)
+            except ModuleNotFoundError:
+                logging.warning("Cannot use sqlite_vec for vector search")
         self.conn.execute(
             f"CREATE TABLE IF NOT EXISTS {self.name} (entity_key BLOB, feature_name TEXT, value BLOB, vector_value BLOB, event_ts timestamp, created_ts timestamp,  PRIMARY KEY(entity_key, feature_name))"
         )
