fix pop

Signed-off-by: Praneeth Bedapudi <praneeth@bpraneeth.com>
notAI-tech · Nov 6, 2023 · 894398a · 894398a
1 parent 99a86ab
commit 894398a
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 29 deletions.
diff --git a/liteindex/defined_index.py b/liteindex/defined_index.py
@@ -381,7 +381,7 @@ def search(
         n=None,
         page_no=None,
         select_keys=[],
-        update=None,
+        update=None
     ):
         if {k for k in query if k not in self.schema or self.schema[k] in {"other"}}:
             raise ValueError("Invalid query")
@@ -475,7 +475,18 @@ def group(self, keys, query):
         }
 
     def pop(self, ids=None, query={}, n=1, sort_by=None, reversed_sort=False):
-        if query:
+        if ids is not None:
+            return {
+                row[0]: self.deserialize_record(
+                    {h: val for h, val in zip(self.column_names, row[2:]) if h in self.key_hash_to_original_key}
+                )
+                for row in self._connection.execute(
+                    f"DELETE FROM {self.name} WHERE id IN ({', '.join(['?' for _ in ids])}) RETURNING *",
+                    ids,
+                ).fetchall()
+            }
+
+        elif query is not None:
             sql_query, sql_params = pop_query(
                 table_name=self.name,
                 query={self.original_key_to_key_hash[k]: v for k, v in query.items()},
@@ -485,16 +496,12 @@ def pop(self, ids=None, query={}, n=1, sort_by=None, reversed_sort=False):
                 n=n,
             )
 
-            return [
-                {
-                    self.key_hash_to_original_key[h]: val
-                    for h, val in zip(self.column_names, row[1:])
-                }
+            return {
+                row[0]: self.deserialize_record(
+                    {h: val for h, val in zip(self.column_names, row[2:]) if h in self.key_hash_to_original_key}
+                )
                 for row in self._connection.execute(sql_query, sql_params).fetchall()
-            ]
-
-        elif ids:
-            pass
+            }
 
         else:
             raise ValueError("Either ids or query must be provided")

diff --git a/liteindex/query_parser.py b/liteindex/query_parser.py
@@ -94,35 +94,26 @@ def pop_query(
     schema,
     sort_by=None,
     reversed_sort=False,
-    n=None,
+    n=None
 ):
     # Prepare the query
     where_conditions, params = parse_query(query, schema)
 
     # Build the query string
     query_str = f"DELETE FROM {table_name}"
     if where_conditions:
-        query_str += f" WHERE {' AND '.join(where_conditions)}"
+        query_str += f" WHERE id IN (SELECT id FROM {table_name} WHERE {' AND '.join(where_conditions)}"
+    else:
+        query_str += f" WHERE id IN (SELECT id FROM {table_name}"
+
     if sort_by:
-        if isinstance(sort_by, list):
-            query_str += " ORDER BY "
-            sort_list = []
-            for sort_item in sort_by:
-                if isinstance(sort_item, tuple):
-                    sort_list.append(
-                        f"{sort_item[0]} {'DESC' if sort_item[1] else 'ASC'}"
-                    )
-                else:
-                    sort_list.append(
-                        f"{sort_item} {'DESC' if reversed_sort else 'ASC'}"
-                    )
-            query_str += ", ".join(sort_list)
-        else:
-            query_str += f" ORDER BY {sort_by} {'DESC' if reversed_sort else 'ASC'}"
+        query_str += f" ORDER BY {sort_by} {'DESC' if reversed_sort else 'ASC'}"
 
     if n is not None:
         query_str += f" LIMIT {n}"
-
+
+    query_str += ")"
+
     query_str += " RETURNING *"
 
     return query_str, params