miguelgrinberg · Jul 30, 2024
diff --git a/app/__init__.py b/app/__init__.py
@@ -8,6 +8,7 @@
 from flask_mail import Mail
 from flask_moment import Moment
 from flask_babel import Babel, lazy_gettext as _l
+from elasticsearch import Elasticsearch
 from config import Config
 
 
@@ -35,6 +36,8 @@ def create_app(config_class=Config):
     mail.init_app(app)
     moment.init_app(app)
     babel.init_app(app, locale_selector=get_locale)
+    app.elasticsearch = Elasticsearch([app.config['ELASTICSEARCH_URL']]) \
+        if app.config['ELASTICSEARCH_URL'] else None
 
     from app.errors import bp as errors_bp
     app.register_blueprint(errors_bp)

diff --git a/app/main/forms.py b/app/main/forms.py
@@ -1,3 +1,4 @@
+from flask import request
 from flask_wtf import FlaskForm
 from wtforms import StringField, SubmitField, TextAreaField
 from wtforms.validators import ValidationError, DataRequired, Length
@@ -33,3 +34,14 @@ class PostForm(FlaskForm):
     post = TextAreaField(_l('Say something'), validators=[
         DataRequired(), Length(min=1, max=140)])
     submit = SubmitField(_l('Submit'))
+
+
+class SearchForm(FlaskForm):
+    q = StringField(_l('Search'), validators=[DataRequired()])
+
+    def __init__(self, *args, **kwargs):
+        if 'formdata' not in kwargs:
+            kwargs['formdata'] = request.args
+        if 'meta' not in kwargs:
+            kwargs['meta'] = {'csrf': False}
+        super(SearchForm, self).__init__(*args, **kwargs)
diff --git a/app/main/routes.py b/app/main/routes.py
@@ -6,7 +6,7 @@
 import sqlalchemy as sa
 from langdetect import detect, LangDetectException
 from app import db
-from app.main.forms import EditProfileForm, EmptyForm, PostForm
+from app.main.forms import EditProfileForm, EmptyForm, PostForm, SearchForm
 from app.models import User, Post
 from app.translate import translate
 from app.main import bp
@@ -17,6 +17,7 @@ def before_request():
     if current_user.is_authenticated:
         current_user.last_seen = datetime.now(timezone.utc)
         db.session.commit()
+        g.search_form = SearchForm()
     g.locale = str(get_locale())
 
 
@@ -150,3 +151,19 @@ def translate_text():
     return {'text': translate(data['text'],
                               data['source_language'],
                               data['dest_language'])}
+
+
+@bp.route('/search')
+@login_required
+def search():
+    if not g.search_form.validate():
+        return redirect(url_for('main.explore'))
+    page = request.args.get('page', 1, type=int)
+    posts, total = Post.search(g.search_form.q.data, page,
+                               current_app.config['POSTS_PER_PAGE'])
+    next_url = url_for('main.search', q=g.search_form.q.data, page=page + 1) \
+        if total > page * current_app.config['POSTS_PER_PAGE'] else None
+    prev_url = url_for('main.search', q=g.search_form.q.data, page=page - 1) \
+        if page > 1 else None
+    return render_template('search.html', title=_('Search'), posts=posts,
+                           next_url=next_url, prev_url=prev_url)
diff --git a/app/models.py b/app/models.py
@@ -9,6 +9,51 @@
 from werkzeug.security import generate_password_hash, check_password_hash
 import jwt
 from app import db, login
+from app.search import add_to_index, remove_from_index, query_index
+
+
+class SearchableMixin:
+    @classmethod
+    def search(cls, expression, page, per_page):
+        ids, total = query_index(cls.__tablename__, expression, page, per_page)
+        if total == 0:
+            return [], 0
+        when = []
+        for i in range(len(ids)):
+            when.append((ids[i], i))
+        query = sa.select(cls).where(cls.id.in_(ids)).order_by(
+            db.case(*when, value=cls.id))
+        return db.session.scalars(query), total
+
+    @classmethod
+    def before_commit(cls, session):
+        session._changes = {
+            'add': list(session.new),
+            'update': list(session.dirty),
+            'delete': list(session.deleted)
+        }
+
+    @classmethod
+    def after_commit(cls, session):
+        for obj in session._changes['add']:
+            if isinstance(obj, SearchableMixin):
+                add_to_index(obj.__tablename__, obj)
+        for obj in session._changes['update']:
+            if isinstance(obj, SearchableMixin):
+                add_to_index(obj.__tablename__, obj)
+        for obj in session._changes['delete']:
+            if isinstance(obj, SearchableMixin):
+                remove_from_index(obj.__tablename__, obj)
+        session._changes = None
+
+    @classmethod
+    def reindex(cls):
+        for obj in db.session.scalars(sa.select(cls)):
+            add_to_index(cls.__tablename__, obj)
+
+
+db.event.listen(db.session, 'before_commit', SearchableMixin.before_commit)
+db.event.listen(db.session, 'after_commit', SearchableMixin.after_commit)
 
 
 followers = sa.Table(
@@ -113,7 +158,8 @@ def load_user(id):
     return db.session.get(User, int(id))
 
 
-class Post(db.Model):
+class Post(SearchableMixin, db.Model):
+    __searchable__ = ['body']
     id: so.Mapped[int] = so.mapped_column(primary_key=True)
     body: so.Mapped[str] = so.mapped_column(sa.String(140))
     timestamp: so.Mapped[datetime] = so.mapped_column(

diff --git a/app/search.py b/app/search.py
@@ -0,0 +1,28 @@
+from flask import current_app
+
+
+def add_to_index(index, model):
+    if not current_app.elasticsearch:
+        return
+    payload = {}
+    for field in model.__searchable__:
+        payload[field] = getattr(model, field)
+    current_app.elasticsearch.index(index=index, id=model.id, document=payload)
+
+
+def remove_from_index(index, model):
+    if not current_app.elasticsearch:
+        return
+    current_app.elasticsearch.delete(index=index, id=model.id)
+
+
+def query_index(index, query, page, per_page):
+    if not current_app.elasticsearch:
+        return [], 0
+    search = current_app.elasticsearch.search(
+        index=index,
+        query={'multi_match': {'query': query, 'fields': ['*']}},
+        from_=(page - 1) * per_page,
+        size=per_page)
+    ids = [int(hit['_id']) for hit in search['hits']['hits']]
+    return ids, search['hits']['total']['value']
diff --git a/app/templates/base.html b/app/templates/base.html
@@ -29,6 +29,13 @@
             <li class="nav-item">
               <a class="nav-link" aria-current="page" href="{{ url_for('main.explore') }}">{{ _('Explore') }}</a>
             </li>
+            {% if g.search_form %}
+            <form class="navbar-form navbar-left" method="get" action="{{ url_for('main.search') }}">
+                <div class="form-group">
+                    {{ g.search_form.q(size=20, class='form-control', placeholder=g.search_form.q.label.text) }}
+                </div>
+            </form>
+            {% endif %}
           </ul>
           <ul class="navbar-nav mb-2 mb-lg-0">
             {% if current_user.is_anonymous %}

diff --git a/app/templates/search.html b/app/templates/search.html
@@ -0,0 +1,22 @@
+{% extends "base.html" %}
+
+{% block content %}
+    <h1>{{ _('Search Results') }}</h1>
+    {% for post in posts %}
+        {% include '_post.html' %}
+    {% endfor %}
+    <nav aria-label="Post navigation">
+        <ul class="pagination">
+            <li class="page-item{% if not prev_url %} disabled{% endif %}">
+                <a class="page-link" href="{{ prev_url }}">
+                    <span aria-hidden="true">&larr;</span> {{ _('Newer posts') }}
+                </a>
+            </li>
+            <li class="page-item{% if not next_url %} disabled{% endif %}">
+                <a class="page-link" href="{{ next_url }}">
+                    {{ _('Older posts') }} <span aria-hidden="true">&rarr;</span>
+                </a>
+            </li>
+        </ul>
+    </nav>
+{% endblock %}
diff --git a/app/translations/es/LC_MESSAGES/messages.po b/app/translations/es/LC_MESSAGES/messages.po
@@ -7,7 +7,7 @@ msgid ""
 msgstr ""
 "Project-Id-Version: PROJECT VERSION\n"
 "Report-Msgid-Bugs-To: EMAIL@ADDRESS\n"
-"POT-Creation-Date: 2017-11-25 17:17-0800\n"
+"POT-Creation-Date: 2017-11-25 18:23-0800\n"
 "PO-Revision-Date: 2017-09-29 23:25-0700\n"
 "Last-Translator: FULL NAME <EMAIL@ADDRESS>\n"
 "Language: es\n"
@@ -18,7 +18,7 @@ msgstr ""
 "Content-Transfer-Encoding: 8bit\n"
 "Generated-By: Babel 2.5.1\n"
 
-#: app/__init__.py:17
+#: app/__init__.py:18
 msgid "Please log in to access this page."
 msgstr "Por favor ingrese para acceder a esta página."
 
@@ -34,43 +34,43 @@ msgstr "Error el servicio de traducciones ha fallado."
 msgid "[Microblog] Reset Your Password"
 msgstr "[Microblog] Nueva Contraseña"
 
-#: app/auth/forms.py:9 app/auth/forms.py:16 app/main/forms.py:10
+#: app/auth/forms.py:10 app/auth/forms.py:17 app/main/forms.py:10
 msgid "Username"
 msgstr "Nombre de usuario"
 
-#: app/auth/forms.py:10 app/auth/forms.py:18 app/auth/forms.py:41
+#: app/auth/forms.py:11 app/auth/forms.py:19 app/auth/forms.py:42
 msgid "Password"
 msgstr "Contraseña"
 
-#: app/auth/forms.py:11
+#: app/auth/forms.py:12
 msgid "Remember Me"
 msgstr "Recordarme"
 
-#: app/auth/forms.py:12 app/templates/auth/login.html:5
+#: app/auth/forms.py:13 app/templates/auth/login.html:5
 msgid "Sign In"
 msgstr "Ingresar"
 
-#: app/auth/forms.py:17 app/auth/forms.py:36
+#: app/auth/forms.py:18 app/auth/forms.py:37
 msgid "Email"
 msgstr "Email"
 
-#: app/auth/forms.py:20 app/auth/forms.py:43
+#: app/auth/forms.py:21 app/auth/forms.py:44
 msgid "Repeat Password"
 msgstr "Repetir Contraseña"
 
-#: app/auth/forms.py:22 app/templates/auth/register.html:5
+#: app/auth/forms.py:23 app/templates/auth/register.html:5
 msgid "Register"
 msgstr "Registrarse"
 
-#: app/auth/forms.py:27 app/main/forms.py:23
+#: app/auth/forms.py:28 app/main/forms.py:23
 msgid "Please use a different username."
 msgstr "Por favor use un nombre de usuario diferente."
 
-#: app/auth/forms.py:32
+#: app/auth/forms.py:33
 msgid "Please use a different email address."
 msgstr "Por favor use una dirección de email diferente."
 
-#: app/auth/forms.py:37 app/auth/forms.py:45
+#: app/auth/forms.py:38 app/auth/forms.py:46
 msgid "Request Password Reset"
 msgstr "Pedir una nueva contraseña"
 
@@ -102,37 +102,41 @@ msgstr "Enviar"
 msgid "Say something"
 msgstr "Dí algo"
 
-#: app/main/routes.py:35
+#: app/main/forms.py:32
+msgid "Search"
+msgstr "Buscar"
+
+#: app/main/routes.py:36
 msgid "Your post is now live!"
 msgstr "¡Tu artículo ha sido publicado!"
 
-#: app/main/routes.py:86
+#: app/main/routes.py:87
 msgid "Your changes have been saved."
 msgstr "Tus cambios han sido salvados."
 
-#: app/main/routes.py:91 app/templates/edit_profile.html:5
+#: app/main/routes.py:92 app/templates/edit_profile.html:5
 msgid "Edit Profile"
 msgstr "Editar Perfil"
 
-#: app/main/routes.py:100 app/main/routes.py:116
+#: app/main/routes.py:101 app/main/routes.py:117
 #, python-format
 msgid "User %(username)s not found."
 msgstr "El usuario %(username)s no ha sido encontrado."
 
-#: app/main/routes.py:103
+#: app/main/routes.py:104
 msgid "You cannot follow yourself!"
 msgstr "¡No te puedes seguir a tí mismo!"
 
-#: app/main/routes.py:107
+#: app/main/routes.py:108
 #, python-format
 msgid "You are following %(username)s!"
 msgstr "¡Ahora estás siguiendo a %(username)s!"
 
-#: app/main/routes.py:119
+#: app/main/routes.py:120
 msgid "You cannot unfollow yourself!"
 msgstr "¡No te puedes dejar de seguir a tí mismo!"
 
-#: app/main/routes.py:123
+#: app/main/routes.py:124
 #, python-format
 msgid "You are not following %(username)s."
 msgstr "No estás siguiendo a %(username)s."
@@ -158,19 +162,19 @@ msgstr "Inicio"
 msgid "Explore"
 msgstr "Explorar"
 
-#: app/templates/base.html:26
+#: app/templates/base.html:33
 msgid "Login"
 msgstr "Ingresar"
 
-#: app/templates/base.html:28
+#: app/templates/base.html:35
 msgid "Profile"
 msgstr "Perfil"
 
-#: app/templates/base.html:29
+#: app/templates/base.html:36
 msgid "Logout"
 msgstr "Salir"
 
-#: app/templates/base.html:66
+#: app/templates/base.html:73
 msgid "Error: Could not contact server."
 msgstr "Error: el servidor no pudo ser contactado."
 
@@ -187,6 +191,18 @@ msgstr "Artículos siguientes"
 msgid "Older posts"
 msgstr "Artículos previos"
 
+#: app/templates/search.html:4
+msgid "Search Results"
+msgstr "Resultados de Búsqueda"
+
+#: app/templates/search.html:12
+msgid "Previous results"
+msgstr "Resultados previos"
+
+#: app/templates/search.html:17
+msgid "Next results"
+msgstr "Resultados próximos"
+
 #: app/templates/user.html:8
 msgid "User"
 msgstr "Usuario"
@@ -256,3 +272,4 @@ msgstr "Ha ocurrido un error inesperado"
 #: app/templates/errors/500.html:5
 msgid "The administrator has been notified. Sorry for the inconvenience!"
 msgstr "El administrador ha sido notificado. ¡Lamentamos la inconveniencia!"
+
diff --git a/config.py b/config.py
@@ -17,4 +17,5 @@ class Config:
     ADMINS = ['your-email@example.com']
     LANGUAGES = ['en', 'es']
     MS_TRANSLATOR_KEY = os.environ.get('MS_TRANSLATOR_KEY')
+    ELASTICSEARCH_URL = os.environ.get('ELASTICSEARCH_URL')
     POSTS_PER_PAGE = 25
diff --git a/requirements.txt b/requirements.txt
@@ -8,6 +8,8 @@ certifi==2023.11.17
 charset-normalizer==3.3.2
 click==8.1.7
 dnspython==2.4.2
+elastic-transport==8.10.0
+elasticsearch==8.11.0
 email-validator==2.1.0.post1
 Flask==3.0.0
 flask-babel==4.0.0

diff --git a/tests.py b/tests.py
@@ -9,6 +9,7 @@
 class TestConfig(Config):
     TESTING = True
     SQLALCHEMY_DATABASE_URI = 'sqlite://'
+    ELASTICSEARCH_URL = None
 
 
 class UserModelCase(unittest.TestCase):