Skip to content

Commit

Permalink
fix: init_es.py script not completing due to TypeError (pinterest#1309)
Browse files Browse the repository at this point in the history
  • Loading branch information
vinmh authored and aidenprice committed Jan 3, 2024
1 parent 8a75a25 commit 0a6a63b
Showing 1 changed file with 8 additions and 7 deletions.
15 changes: 8 additions & 7 deletions querybook/server/logic/elasticsearch.py
Original file line number Diff line number Diff line change
Expand Up @@ -697,7 +697,7 @@ def get_fullname_field():


@with_session
def get_users_iter(batch_size=5000, fields=None, session=None):
def get_users_iter(batch_size=5000, session=None):
offset = 0

while True:
Expand All @@ -709,8 +709,7 @@ def get_users_iter(batch_size=5000, fields=None, session=None):
LOG.info("\n--User count: {}, offset: {}".format(len(users), offset))

for user in users:
expanded_user = user_to_es(user, fields=fields, session=session)
yield expanded_user
yield user

if len(users) < batch_size:
break
Expand All @@ -723,15 +722,17 @@ def _bulk_insert_users():
for user in get_users_iter():
# skip indexing user groups before having the correct permission setup for it.
if not user.is_group:
_insert(index_name, user["id"], user)
expanded_user = user_to_es(user, fields=None)
_insert(index_name, expanded_user["id"], expanded_user)


def _bulk_update_users(fields: Set[str] = None):
index_name = ES_CONFIG["users"]["index_name"]

for user in get_users_iter(fields=fields):
updated_body = {"doc": user, "doc_as_upsert": True}
_update(index_name, user["id"], updated_body)
for user in get_users_iter():
expanded_user = user_to_es(user, fields=fields)
updated_body = {"doc": expanded_user, "doc_as_upsert": True}
_update(index_name, expanded_user["id"], updated_body)


@with_exception
Expand Down

0 comments on commit 0a6a63b

Please sign in to comment.