alikia2x/OpenRewind
1
1
Fork 0
Code Issues Releases 1 Activity

Improved readability

Browse Source
This commit is contained in:
Koen van Eijk 2024-06-09 01:06:21 +02:00
parent be9c7de2b1
commit 1151e5f0af
1 changed files with 35 additions and 51 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

54
openrecall/database.py
View File

@ -1,62 +1,46 @@
import sqlite3 import sqlite3
from collections import namedtuple
from typing import Any, List
from openrecall.config import db_path from openrecall.config import db_path
Entry = namedtuple("Entry", ["id", "app", "title", "text", "timestamp", "embedding"])
def create_db():
conn = sqlite3.connect(db_path) def create_db() -> None:
with sqlite3.connect(db_path) as conn:
c = conn.cursor() c = conn.cursor()
c.execute( c.execute(
"""CREATE TABLE IF NOT EXISTS entries """CREATE TABLE IF NOT EXISTS entries
(id INTEGER PRIMARY KEY AUTOINCREMENT, app TEXT, title TEXT, text TEXT, timestamp INTEGER, embedding BLOB)""" (id INTEGER PRIMARY KEY AUTOINCREMENT, app TEXT, title TEXT, text TEXT, timestamp INTEGER, embedding BLOB)"""
) )
conn.commit() conn.commit()
conn.close()
def get_all_entries(): def get_all_entries() -> List[Entry]:
conn = sqlite3.connect(db_path) with sqlite3.connect(db_path) as conn:
c = conn.cursor() c = conn.cursor()
results = c.execute("SELECT * FROM entries").fetchall() results = c.execute("SELECT * FROM entries").fetchall()
conn.close() return [Entry(*result) for result in results]
entries = []
for result in results:
entries.append(
{
"id": result[0],
"app": result[1],
"title": result[2],
"text": result[3],
"timestamp": result[4],
"embedding": result[5],
}
)
return entries
def get_timestamps(): def get_timestamps() -> List[int]:
conn = sqlite3.connect(db_path) with sqlite3.connect(db_path) as conn:
c = conn.cursor() c = conn.cursor()
results = c.execute( results = c.execute(
"SELECT timestamp FROM entries ORDER BY timestamp DESC LIMIT 1000" "SELECT timestamp FROM entries ORDER BY timestamp DESC LIMIT 1000"
).fetchall() ).fetchall()
timestamps = [result[0] for result in results] return [result[0] for result in results]
conn.close()
return timestamps
def insert_entry(text, timestamp, embedding, app, title):
conn = sqlite3.connect(db_path) def insert_entry(
c = conn.cursor() text: str, timestamp: int, embedding: Any, app: str, title: str
) -> None:
embedding_bytes = embedding.tobytes() embedding_bytes = embedding.tobytes()
with sqlite3.connect(db_path) as conn:
c = conn.cursor()
c.execute( c.execute(
"INSERT INTO entries (text, timestamp, embedding, app, title) VALUES (?, ?, ?, ?, ?)", "INSERT INTO entries (text, timestamp, embedding, app, title) VALUES (?, ?, ?, ?, ?)",
( (text, timestamp, embedding_bytes, app, title),
text,
timestamp,
embedding_bytes,
app,
title,
),
) )
conn.commit() conn.commit()
conn.close()