mysql -> sqlite

Author: davidheineman

.gitignore

@@ -5,4 +5,5 @@ app
 index
 *.json
 *.bib
-main.yml
+main.yml
+anthology.db

docker-compose.yml

@@ -3,6 +3,6 @@ colbert-ir[torch,faiss-cpu] # faiss-gpu
 faiss-cpu
 huggingface_hub
 bibtexparser
-mysql-connector-python
+pysqlite3
 flask
 gunicorn

requirements.txt

@@ -11,7 +11,6 @@
 DATASET_PATH = os.path.join(INDEX_ROOT, 'dataset.json')
 
 DB_NAME = 'anthology'
-DB_HOSTNAME = 'mysql_db' # localhost
-DB_PORT = 3306 # None
+DB_FILENAME = 'anthology.db'
 
 VENUES = Literal['workshop', 'journal', 'short', 'demo', 'tutorial', 'industry', 'findings', 'main']

src/constants.py

@@ -1,8 +1,8 @@
 import json
 from typing import List, Optional, Union
 
-import mysql.connector
-from constants import DATASET_PATH, DB_HOSTNAME, DB_NAME, DB_PORT, VENUES
+import sqlite3
+from constants import DATASET_PATH, DB_FILENAME, DB_NAME, VENUES
 
 PAPER_QUERY = """
 SELECT * 
@@ -19,27 +19,11 @@ def read_dataset():
 
 
 def create_database():
-    db = mysql.connector.connect(
-        host = DB_HOSTNAME,
-        user = "root",
-        password = "",
-        port = DB_PORT
+    db: sqlite3.Connection = sqlite3.connect(
+        database = DB_FILENAME
     )
     cursor = db.cursor()
 
-    cursor.execute("SHOW DATABASES")
-    db_exists = False
-    for x in cursor:
-        db_name = x[0]
-        if db_name == DB_NAME:
-            db_exists = True
-
-    # Create database
-    if not db_exists:
-        print("Creating new database...")
-        cursor.execute(f'CREATE DATABASE {DB_NAME}')
-    cursor.execute(f'USE {DB_NAME}')
-
     # Create table
     print('Creating new table...')
     cursor.execute(f'DROP TABLE IF EXISTS paper')
@@ -61,6 +45,7 @@ def create_database():
     acl_data = read_dataset()
 
     vals = []
+    paper: dict
     for pid, paper in enumerate(acl_data):
         title    = paper.get('title', '')
         author   = paper.get('author', '')
@@ -80,7 +65,7 @@ def create_database():
     INSERT INTO paper (
         pid, title, author, year, abstract, url, type, venue, venue_type, is_findings
     ) VALUES (
-        %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
+        ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
     )
     """
 
@@ -123,17 +108,13 @@ def query_paper_metadata(
     ):
     if not isinstance(venue_type, list): venue_type = [venue_type]
 
-    db = mysql.connector.connect(
-        host = DB_HOSTNAME,
-        user = "root",
-        password = "",
-        database = DB_NAME,
-        port = DB_PORT
+    db: sqlite3.Connection = sqlite3.connect(
+        database = DB_FILENAME
     )
 
     cursor = db.cursor()
 
-    pids_str = ', '.join(['%s'] * len(pids))
+    pids_str = ', '.join(['?'] * len(pids))
 
     constraints_str = ""
     if start_year: constraints_str += f" AND year >= {start_year}"

src/db.py

@@ -69,9 +69,9 @@ def query():
     # Get top passage IDs from ColBERT
     colbert_response = api_search_query(query)
 
-    # Query MySQL database for paper information
+    # Query database for paper information
     pids = [r['pid'] for r in colbert_response["topk"]]
-    mysql_response = query_paper_metadata(
+    server_response = query_paper_metadata(
         pids, 
         start_year=start_year,
         end_year=end_year,
@@ -80,9 +80,9 @@ def query():
     )
 
     K = 20
-    mysql_response = mysql_response[:K]
+    server_response = server_response[:K]
 
-    return mysql_response
+    return server_response
 
 
 # @app.route('/search', methods=['POST', 'GET'])