1 files changed, 9 insertions, 6 deletions
diff --git a/check/app/models/sql_factory.py b/check/app/models/sql_factory.py
index 3ed3af0..d4a371e 100644
--- a/check/app/models/sql_factory.py
+++ b/check/app/models/sql_factory.py
@@ -3,7 +3,7 @@ import glob
 import time
 import pandas as pd
 
-from sqlalchemy import create_engine, Table, Column, String, BigInteger
+from sqlalchemy import create_engine, Table, Column, String, Integer, BigInteger
 from sqlalchemy.orm import sessionmaker
 from sqlalchemy.ext.declarative import declarative_base
 
@@ -22,11 +22,12 @@ Session = sessionmaker(bind=engine)
 
 Base = declarative_base()
 
-class PHashTable(Base):
-  __tablename__ = 'phashes'
+class FileTable(Base):
+  """Table for storing various hashes of images"""
+  __tablename__ = 'files'
   id = Column(Integer, primary_key=True)
   sha256 = Column(String(36, convert_unicode=True), nullable=False)
-  phash = Column(BigInteger(blank=True), nullable=False)
+  phash = Column(BigInteger, nullable=False)
   ext = Column(String(3, convert_unicode=True), nullable=False)
   def toJSON(self):
     return {
@@ -37,13 +38,15 @@ class PHashTable(Base):
     }
 
 def search_by_phash(phash, threshold=6):
+  """Search files for a particular phash"""
   connection = engine.connect()
-  cmd = 'SELECT phashes.*, BIT_COUNT(phash ^ :phash) as hamming_distance FROM images_image HAVING hamming_distance < :threshold ORDER BY hamming_distance ASC LIMIT 1'
+  cmd = 'SELECT files.*, BIT_COUNT(phash ^ :phash) as hamming_distance FROM images_image HAVING hamming_distance < :threshold ORDER BY hamming_distance ASC LIMIT 1'
   matches = connection.execute(text(cmd), phash=phash, threshold=threshold)
   return matches
 
 def add_phash(sha256, phash, ext):
-  rec = PHashTable(
+  """Add a file to the table"""
+  rec = FileTable(
     sha256=sha256, phash=phash, ext=ext,
   )
   session = Session()