Spaces:
Sleeping
Sleeping
File size: 1,893 Bytes
cb90daa f543e54 cb90daa f543e54 cb90daa f543e54 cb90daa f550818 cb90daa f543e54 cb90daa |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import sqlite3
import json
class create_db:
def __init__(self, token, chunk_json1, filename):
conn = sqlite3.connect('my_database.db')
cursor = conn.cursor()
cursor.execute("""
CREATE TABLE IF NOT EXISTS token_data (
token_id TEXT PRIMARY KEY,
chunk_data TEXT,
filename TEXT
)
""")
chunk_json = json.dumps(chunk_json1)
try:
cursor.execute(
"INSERT INTO token_data (token_id, chunk_data, filename) VALUES (?, ?, ?)",
(token, chunk_json, filename)
)
conn.commit()
print({"message": f"✅ {filename} uploaded and stored successfully"})
except sqlite3.IntegrityError:
print({"error": f"❌ Token already exists for: {filename}"})
conn.close()
@staticmethod
def get_all_filenames():
conn = sqlite3.connect('my_database.db')
cursor = conn.cursor()
cursor.execute("SELECT filename FROM token_data")
rows = cursor.fetchall()
conn.close()
return {"pdfs": [{"filename": row[0]} for row in rows]}
# ✅ NEW: Fetch record by filename and token
@staticmethod
def fetch_by_token_or_filename(identifier):
conn = sqlite3.connect('my_database.db')
cursor = conn.cursor()
cursor.execute("""
SELECT token_id, chunk_data, filename
FROM token_data
WHERE token_id = ? OR filename = ?
""", (identifier, identifier))
result = cursor.fetchone()
conn.close()
if result:
return {
"token": result[0],
"chunk_data": result[1],
"filename": result[2],
}
else:
return {"error": "No matching record found for token or filename."} |