update json querying

Signed-off-by: Praneeth Bedapudi <praneeth@bpraneeth.com>
notAI-tech · Oct 11, 2024 · 8c3bbba · 8c3bbba
1 parent fa7f1c7
commit 8c3bbba
Show file tree

Hide file tree

Showing 4 changed files with 68 additions and 36 deletions.
diff --git a/liteindex/kv_index.py b/liteindex/kv_index.py
@@ -495,22 +495,18 @@ def search(
         return results
 
     def math(self, key, value, op):
-        ops = {
-            "+=": "+",
-            "-=": "-",
-            "*=": "*",
-            "/=": "/",
-            "%=": "%"
-        }
+        ops = {"+=": "+", "-=": "-", "*=": "*", "/=": "/", "%=": "%"}
 
         if op not in ops:
             raise ValueError(f"Unsupported operation: {op}")
 
         sql_query = f"UPDATE kv_index SET num_value = num_value {ops[op]} ? WHERE key_hash = ? RETURNING num_value"
 
         with self.__connection as conn:
-            row = conn.execute(sql_query, (value, self.__encode_and_hash(key)[0])).fetchone()
-
+            row = conn.execute(
+                sql_query, (value, self.__encode_and_hash(key)[0])
+            ).fetchone()
+
             if row is None:
                 raise KeyError
 

diff --git a/liteindex/query_parser.py b/liteindex/query_parser.py
@@ -43,10 +43,9 @@ def process_nested_query(value, prefix):
                         "$lte": "<=",
                     }[sub_key]
 
-                    # Handle the case for JSON columns with "$like" operator
                     if schema[prefix[0]] == "json" and sub_key == "$like":
                         sub_conditions.append(
-                            f"JSON_EXTRACT({column}, '$[*]') {operator} ?"
+                            f"EXISTS(SELECT 1 FROM json_each({column}) WHERE value {operator} ?)"
                         )
                     else:
                         if sub_key == "$ne":
@@ -64,41 +63,70 @@ def process_nested_query(value, prefix):
 
                     params.append(sub_value)
                 elif sub_key == "$in":
-                    sub_conditions.append(
-                        f"""("{column}" IN ({', '.join(['?' for _ in sub_value])}) OR "{column}" IS NULL)"""
-                        if column_type is not None
-                        else f"({column} IN ({', '.join(['?' for _ in sub_value])}) OR {column} IS NULL)"
-                    )
-                    params.extend(sub_value)
+                    if schema[prefix[0]] == "json":
+                        json_conditions = []
+                        for val in sub_value:
+                            json_conditions.append(
+                                f"EXISTS(SELECT 1 FROM json_each({column}) WHERE value = ?)"
+                            )
+                            params.append(val)
+                        sub_conditions.append(f"({ ' OR '.join(json_conditions) })")
+                    else:
+                        sub_conditions.append(
+                            f"""("{column}" IN ({', '.join(['?' for _ in sub_value])}) OR "{column}" IS NULL)"""
+                            if column_type is not None
+                            else f"({column} IN ({', '.join(['?' for _ in sub_value])}) OR {column} IS NULL)"
+                        )
+                        params.extend(sub_value)
                 elif sub_key == "$nin":
-                    sub_conditions.append(
-                        f"""("{column}" NOT IN ({', '.join(['?' for _ in sub_value])}) OR "{column}" IS NULL)"""
-                        if column_type is not None
-                        else f"({column} NOT IN ({', '.join(['?' for _ in sub_value])}) OR {column} IS NULL)"
-                    )
-                    params.extend(sub_value)
+                    if schema[prefix[0]] == "json":
+                        json_conditions = []
+                        for val in sub_value:
+                            json_conditions.append(
+                                f"NOT EXISTS(SELECT 1 FROM json_each({column}) WHERE value = ?)"
+                            )
+                            params.append(val)
+                        sub_conditions.append(f"({ ' AND '.join(json_conditions) })")
+                    else:
+                        sub_conditions.append(
+                            f"""("{column}" NOT IN ({', '.join(['?' for _ in sub_value])}) OR "{column}" IS NULL)"""
+                            if column_type is not None
+                            else f"({column} NOT IN ({', '.join(['?' for _ in sub_value])}) OR {column} IS NULL)"
+                        )
+                        params.extend(sub_value)
                 else:
                     process_nested_query(sub_value, prefix + [sub_key])
             if sub_conditions:
                 where_conditions.append(f"({' AND '.join(sub_conditions)})")
 
         elif isinstance(value, list):
             if schema[prefix[0]] == "json":
-                json_conditions = [f"JSON_CONTAINS({column}, ?)" for _ in value]
+                json_conditions = []
+                for val in value:
+                    json_conditions.append(
+                        f"EXISTS(SELECT 1 FROM json_each({column}) WHERE value = ?)"
+                    )
+                    params.append(json.dumps(val))
                 where_conditions.append(f"({ ' OR '.join(json_conditions) })")
-                params.extend(json.dumps(val) for val in value)
             else:
                 where_conditions.append(
                     f"""("{column}" IN ({', '.join(['?' for _ in value])}) OR "{column}" IS NULL)"""
                 )
                 params.extend(value)
 
         elif value is None:
-            where_conditions.append(
-                f'"{column}" IS NULL'
-                if column_type is not None
-                else f"{column} IS NULL"
-            )
+            if len(prefix) > 1:
+                json_column = prefix[0]
+                json_path = ".".join(prefix[1:])
+                where_conditions.append(
+                    f"""(json_type({json_column}, '$.{json_path}') IS NOT NULL AND json_extract({json_column}, '$.{json_path}') IS NULL)"""
+                )
+            else:
+                where_conditions.append(
+                    f'"{column}" IS NULL'
+                    if column_type is not None
+                    else f"{column} IS NULL"
+                )
         else:
             if column_type == "other":
                 column = f"__hash_{column}"

diff --git a/setup.py b/setup.py
@@ -18,7 +18,7 @@
 EMAIL = "praneeth@bpraneeth.com"
 AUTHOR = "BEDAPUDI PRANEETH"
 REQUIRES_PYTHON = ">=3.6.0"
-VERSION = "0.0.2.dev64"
+VERSION = "0.0.3"
 
 # What packages are required for this module to be executed?
 REQUIRED = []

diff --git a/tests/test_json.py b/tests/test_json.py
@@ -49,8 +49,16 @@
     }
 )
 
-print(index.search({"json_key": {"a": 1}}))
-print(index.search({"json_key": {"a": None}}))
-print(index.search({"json_key": {"a": {"$gte": 1}}}))
-print(index.search({"json_key": {"a": {"$in": [-1, 99]}}}))
-# print(index.search({"json_key": ["c"]}))
+# print(index.search({"json_key": {"$in": ["a", "c"]}}))
+
+assert list(index.search({"json_key": {"a": 1}}).keys()) == ["k1"]
+
+assert list(index.search({"json_key": {"a": None}}).keys()) == ["k4"]
+
+assert list(index.search({"json_key": {"a": {"$gte": 1}}}).keys()) == ["k1", "k2", "k5"]
+
+# assert list(index.search({"json_key": {"a": {"$in": [-1, 99]}}}).keys()) == []
+
+assert list(index.search({"json_key": {"a": {"$ne": 1}}}).keys()) == ["k2", "k3", "k4", "k5"]
+
+assert list(index.search({"json_key": {"a": {"$ne": None}}}).keys()) == ["k1", "k2", "k5"]