1
0
Fork 0
mirror of https://github.com/Findus23/cr-search.git synced 2024-09-19 15:23:44 +02:00
cr-search/models.py

73 lines
1.8 KiB
Python

from peewee import PostgresqlDatabase, Model, IntegerField, CharField, BooleanField, ForeignKeyField
from playhouse.postgres_ext import TSVectorField
from config import dbauth
db = PostgresqlDatabase(**dbauth)
db.connect()
class BaseModel(Model):
class Meta:
database = db
class Episode(BaseModel):
season = IntegerField()
episode_number = IntegerField()
video_number = IntegerField()
youtube_id = CharField(max_length=11)
class Meta:
indexes = ((("season", "video_number"), True),)
@property
def name(self) -> str:
return f"C{self.season}E{self.episode_number:03d}"
class Person(BaseModel):
name = CharField(unique=True)
color = CharField(null=True)
FULL_TEXT_SEARCH = '''SELECT id, text, ts_rank_cd(search_text, query) AS rank
FROM line,
to_tsquery('english', %s) query
WHERE search_text @@ query
ORDER BY rank DESC
'''
class Line(BaseModel):
text = CharField()
search_text = TSVectorField()
person = ForeignKeyField(Person, backref="lines", null=True)
isnote = BooleanField(default=False)
ismeta = BooleanField(default=False)
starttime = IntegerField()
endtime = IntegerField()
episode = ForeignKeyField(Episode, backref="lines")
order = IntegerField()
class Meta:
indexes = ((("episode", "order"), True),)
@classmethod
def full_text_search(cls, query_string: str):
cursor = cls._meta.database.execute_sql(
FULL_TEXT_SEARCH,
(query_string,)
)
result = cursor.fetchall()
cursor.close()
return result
class Phrase(BaseModel):
text = CharField()
count = IntegerField()
episode = ForeignKeyField(Episode, backref="phrase")
class Meta:
indexes = ((("text", "episode"), True),)