Skip to content

Commit

Permalink
Use HTTP basic auth to run queries as wikiuser
Browse files Browse the repository at this point in the history
  • Loading branch information
1stvamp committed Dec 11, 2023
1 parent fc886b6 commit 86bfe12
Showing 1 changed file with 12 additions and 6 deletions.
18 changes: 12 additions & 6 deletions wikipedia/track.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import random
import re
import uuid
from base64 import b64encode
from os import getcwd
from os.path import dirname
from typing import Iterator, List
Expand All @@ -18,6 +19,7 @@

ROLE_IDS = ["managed-role-search-{}".format(uuid.uuid4()) for x in range(0, 500000)]
USER_AUTH = {"username": "wikiuser", "password": "ujd_rbh5quw7GWC@pjc"}
AUTH_BEARER = b64encode(("{}:{}".format(*USER_AUTH.values())).encode('ascii'))
ROLE_TEMPLATE = {
"indices": [
{
Expand Down Expand Up @@ -143,7 +145,7 @@ def params(self):
query = next(self._queries_iterator)
return {
"method": "POST",
"headers": {"es-security-runas-user": USER_AUTH["username"]},
"headers": {"Authorization": "Basic {}".format(AUTH_BEARER)},
"path": f"{SEARCH_APPLICATION_ROOT_ENDPOINT}/{self.search_application_params.name}/_search",
"body": {
"params": {
Expand Down Expand Up @@ -183,6 +185,10 @@ async def create_users_and_roles(es, params):
# For now we'll just work with one user with all the roles
# num_users = params['users']

await es.security.put_user(
username=USER_AUTH["username"], params={"password": USER_AUTH["password"]}
)

await es.indices.refresh(index="wikipedia")
doc_count = await es.count(index="wikipedia")

Expand All @@ -208,17 +214,13 @@ async def create_users_and_roles(es, params):
)

await es.security.put_user(
username=USER_AUTH["username"], params={"roles": ROLE_IDS[0 : num_roles - 1], "password": USER_AUTH["password"]}
username=USER_AUTH["username"], params={"roles": ROLE_IDS[0 : num_roles - 1]}
)
await es.security.put_role(name='impersonator', body={"run_as": ["wikiuser"]}, refresh="wait_for")
await es.security.put_user(username='esbench', params={"roles": ["superuser", "impersonator"]})

await es.indices.refresh(index="wikipedia")


async def reset_indices(es, params):
await es.security.delete_user(username=USER_AUTH["username"])

roles = await es.security.get_role()
for role in roles:
name = list(role.keys())[0]
Expand All @@ -236,6 +238,10 @@ async def reset_indices(es, params):

await es.indices.refresh(index="wikipedia")

await es.security.put_user(
username=USER_AUTH["username"], params={"roles": []}
)


def register(registry):
registry.register_param_source("query-string-search", QueryParamSource)
Expand Down

0 comments on commit 86bfe12

Please sign in to comment.