diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 2dcf00d..86fdc33 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -41,7 +41,9 @@ jobs: - name: Test Python run: | - python3 setup.py install --user test + pip install -e . + pip install httpretty + python -m unittest discover tests check-inclusive-naming: diff --git a/CHANGES.txt b/CHANGES.txt index e64cba8..f6297a0 100644 --- a/CHANGES.txt +++ b/CHANGES.txt @@ -11,3 +11,4 @@ v1.2.5, 2022-07-12 -- Block some more bot useragents v1.2.6, 2022-07-13 -- Block one more useragent - "gh" v1.2.7, 2022-07-15 -- Block more user agents - "Petalbot" v1.3.0, 2023-02-20 -- Add rate limits +v1.4.0, 2-24-06-21 -- Migrate to Flask.Limiter for rate limits diff --git a/README.md b/README.md index 55d8d13..559d63a 100644 --- a/README.md +++ b/README.md @@ -38,11 +38,12 @@ app.add_url_rule( "/docs/search", "docs-search", build_search_view( + app=app session=session, site="maas.io/docs", template_path="docs/search.html", search_engine_id="xxxxxxxxxx", # Optional argument, required by some of our sites - request_limit="500/day", # Allows your to configure the limit at which the user will be forbidden to query more. Defaults to 500 per day + request_limit="500/day", # Allows your to configure the limit at which the user will be forbidden to query more. Defaults to 2000 per day ) ) ``` diff --git a/canonicalwebteam/search/views.py b/canonicalwebteam/search/views.py index 77981fb..93162fb 100644 --- a/canonicalwebteam/search/views.py +++ b/canonicalwebteam/search/views.py @@ -3,26 +3,28 @@ # Packages import flask +from flask_limiter import Limiter +from flask_limiter.util import get_remote_address # Local from canonicalwebteam.search.models import get_search_results -from limits import storage, strategies, parse - -memory_storage = storage.MemoryStorage() -fixed_window = strategies.MovingWindowRateLimiter(memory_storage) class NoAPIKeyError(Exception): pass +limiter = Limiter(get_remote_address) + + def build_search_view( + app, session, site=None, template_path="search.html", search_engine_id="009048213575199080868:i3zoqdwqk8o", site_restricted_search=False, - request_limit="500/day", + request_limit="2000/day", ): """ Build and return a view function that will query the @@ -33,12 +35,13 @@ def build_search_view( from canonicalwebteam.search import build_search_view - app = Flask() + app = Flask(__name__) session = talisker.requests.get_session() app.add_url_rule( "/search", "search", build_search_view( + app, session=session, site="snapcraft.io", template_path="search.html" @@ -46,18 +49,12 @@ def build_search_view( ) """ + limiter.init_app(app) + def search_view(): """ Get search results from Google Custom Search """ - # Rate limit requests to protect from spamming - # To adjust this rate visit - # https://limits.readthedocs.io/en/latest/quickstart.html#examples - limit = parse(request_limit) - rate_limit = fixed_window.hit(limit) - if not rate_limit: - return flask.abort(429, f"The rate limit is: {request_limit}") - # API key should always be provided as an environment variable search_api_key = os.getenv("SEARCH_API_KEY") @@ -72,16 +69,17 @@ def search_view(): results = None if query: - results = get_search_results( - session=session, - api_key=search_api_key, - search_engine_id=search_engine_id, - siteSearch=site_search, - site_restricted_search=site_restricted_search, - query=query, - start=start, - num=num, - ) + with limiter.limit(request_limit): + results = get_search_results( + session=session, + api_key=search_api_key, + search_engine_id=search_engine_id, + siteSearch=site_search, + site_restricted_search=site_restricted_search, + query=query, + start=start, + num=num, + ) return ( flask.render_template( diff --git a/setup.py b/setup.py index 131cbf5..4f9854f 100755 --- a/setup.py +++ b/setup.py @@ -4,7 +4,7 @@ setup( name="canonicalwebteam.search", - version="1.3.0", + version="1.4.0", author="Canonical webteam", author_email="webteam@canonical.com", url="https://github.com/canonical/canonicalwebteam.search", @@ -15,6 +15,10 @@ packages=find_packages(), long_description=open("README.md").read(), long_description_content_type="text/markdown", - install_requires=["Flask>=1.0.2", "user-agents>=2.0.0", "limits>=3.2.0"], + install_requires=[ + "canonicalwebteam.flask-base>=2.0.0", + "user-agents>=2.0.0", + "Flask-Limiter>=3.8.0", + ], tests_require=["httpretty"], ) diff --git a/tests/test_app.py b/tests/test_app.py index 2920e9f..49cbabf 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -48,7 +48,7 @@ def setUp(self): # Default use-case self.app.add_url_rule( - "/search", "search", build_search_view(session=session) + "/search", "search", build_search_view(self.app, session=session) ) # Custom use-case @@ -56,6 +56,7 @@ def setUp(self): "/docs/search", "docs-search", build_search_view( + self.app, session=session, site="maas.io/docs", template_path="docs/search.html", @@ -67,6 +68,7 @@ def setUp(self): "/server/docs/search", "server-docs-search", build_search_view( + self.app, session=session, template_path="docs/search.html", site_restricted_search=True, @@ -78,6 +80,7 @@ def setUp(self): "/server/docs/limited/search", "server-docs-search-limited", build_search_view( + self.app, session=session, template_path="docs/search.html", site_restricted_search=True,