Skip to content

Commit

Permalink
fix: Fix SQLite import issue (#4294)
Browse files Browse the repository at this point in the history
adding try and except block for import

Co-authored-by: Francisco Javier Arceo <franciscojavierarceo@users.noreply.github.com>
  • Loading branch information
franciscojavierarceo and franciscojavierarceo committed Jul 1, 2024
1 parent 7072fd0 commit 398ea3b
Showing 1 changed file with 15 additions and 4 deletions.
19 changes: 15 additions & 4 deletions sdk/python/feast/infra/online_stores/sqlite.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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))"
)
Expand Down

0 comments on commit 398ea3b

Please sign in to comment.