def test_search(args):
print("\n\n==========", args)
found = search_es_full(args)
search_type = args.get("kind", "all")
def print_entity(title, entities):
if not entities:
return
print(f"\n[ {title} ]")
for entity in entities:
print(
" ",
[
entity["user"]["handle"],
entity["user"]["name"],
entity.get("title") or entity.get("playlist_name"),
f"{entity['repost_count']} reposts",
f"{entity['user']['follower_count']} followers",
f"{entity.get('_score')} score",
],
)
def print_users(title, users):
if not users:
return
print(f"\n[ {title} ]")
for user in users:
print(
" ",
[
user["handle"],
user["name"],
f"{user.get('follower_count')} followers",
f"{user.get('is_verified')} verified",
f"{user.get('_score')} score",
],
)
if search_type == "tracks" or search_type == "all":
print_entity("tracks", found["tracks"])
print_entity("saved tracks", found["saved_tracks"])
if search_type == "users" or search_type == "all":
print_users("users", found["users"])
print_users("followed_users", found["followed_users"])
if search_type == "playlists" or search_type == "all":
print_entity("playlists", found["playlists"])
print_entity("saved_playlists", found["saved_playlists"])
if search_type == "albums" or search_type == "all":
print_entity("albums", found["albums"])
print_entity("saved_albums", found["saved_albums"])
def test_get_internal_albums(app_module):
"""Tests we get albums when a user is logged in"""
with app_module.app_context():
db = get_db()
with db.scoped_session() as session:
res = playlist_search_query(session, "album", 10, 0, True, False, 1)
assert len(res["all"]) == 1
assert len(res["saved"]) == 1
search_args = {
"is_auto_complete": True,
"kind": "albums",
"query": "album",
"current_user_id": 1,
"with_users": True,
"limit": 10,
"offset": 0,
"only_downloadable": False,
}
es_res = search_es_full(search_args)
assert len(es_res["albums"]) == 1
assert len(es_res["saved_albums"]) == 1
def test_get_autocomplete_playlists(app_module):
"""Tests we get all tracks with autocomplete"""
with app_module.app_context():
db = get_db()
with db.scoped_session() as session:
res = playlist_search_query(session, "playlist", 10, 0, False, True, None)
assert len(res["all"]) == 1
assert len(res["saved"]) == 0
search_args = {
"is_auto_complete": True,
"kind": "playlists",
"query": "playlist",
"current_user_id": None,
"with_users": True,
"limit": 10,
"offset": 0,
"only_downloadable": False,
}
es_res = search_es_full(search_args)
assert len(es_res["playlists"]) == 1
assert len(es_res["saved_playlists"]) == 0
def test_get_internal_users_no_following(app_module):
"""Tests we get all users for a user that doesn't follow anyone"""
with app_module.app_context():
db = get_db()
with db.scoped_session() as session:
res = user_search_query(session, "user", 10, 0, False, 1)
assert len(res["all"]) == 2
assert len(res["followed"]) == 0
search_args = {
"is_auto_complete": False,
"kind": "users",
"query": "user",
"current_user_id": 1,
"with_users": True,
"limit": 10,
"offset": 0,
"only_downloadable": False,
}
es_res = search_es_full(search_args)
assert len(es_res["users"]) == 2
assert len(es_res["followed_users"]) == 0
def test_get_tracks_external(app_module):
"""Tests we get all tracks, including downloaded"""
with app_module.app_context():
db = get_db()
with db.scoped_session() as session:
res = track_search_query(session, "the track", 10, 0, False, None, False)
assert len(res["all"]) == 2
assert len(res["saved"]) == 0
search_args = {
"is_auto_complete": False,
"kind": "tracks",
"query": "the track",
"current_user_id": None,
"with_users": True,
"limit": 10,
"offset": 0,
"only_downloadable": False,
}
es_res = search_es_full(search_args)
assert len(es_res["tracks"]) == 2
def search(args):
"""Perform a search. `args` should contain `is_auto_complete`,
`query`, `kind`, `current_user_id`, and `only_downloadable`
"""
if os.getenv("audius_elasticsearch_search_enabled"):
try:
resp = search_es_full(args)
return resp
except Exception as e:
logger.error(f"Elasticsearch error: {e}")
search_str = args.get("query")
# when creating query table, we substitute this too
search_str = search_str.replace("&", "and")
kind = args.get("kind", "all")
is_auto_complete = args.get("is_auto_complete")
current_user_id = args.get("current_user_id")
only_downloadable = args.get("only_downloadable")
limit = args.get("limit")
offset = args.get("offset")
searchKind = SearchKind[kind]
results = {}
# Accumulate user_ids for later
user_ids = set()
# Create args for perform_search_query
search_args = {
"search_str": search_str,
"limit": limit,
"offset": offset,
"is_auto_complete": is_auto_complete,
"current_user_id": current_user_id,
"only_downloadable": only_downloadable,
}
if search_str:
db = get_db_read_replica()
# Concurrency approach:
# Spin up a ThreadPoolExecutor for each request to perform_search_query
# to perform the different search types in parallel.
# After each future resolves, we then add users for each entity in a single
# db round trip.
with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
# Keep a mapping of future -> search_type
futures_map = {}
futures = []
# Helper fn to submit a future and add it to bookkeeping data structures
def submit_and_add(search_type):
future = executor.submit(
perform_search_query, db, search_type, search_args
)
futures.append(future)
futures_map[future] = search_type
if searchKind in [SearchKind.all, SearchKind.tracks]:
submit_and_add("tracks")
if searchKind in [SearchKind.all, SearchKind.users]:
submit_and_add("users")
if searchKind in [SearchKind.all, SearchKind.playlists]:
submit_and_add("playlists")
if searchKind in [SearchKind.all, SearchKind.albums]:
submit_and_add("albums")
for future in concurrent.futures.as_completed(futures):
search_result = future.result()
future_type = futures_map[future]
# Add to the final results
# Add to user_ids
if future_type == "tracks":
results["tracks"] = search_result["all"]
results["saved_tracks"] = search_result["saved"]
elif future_type == "users":
results["users"] = search_result["all"]
results["followed_users"] = search_result["followed"]
elif future_type == "playlists":
results["playlists"] = search_result["all"]
results["saved_playlists"] = search_result["saved"]
elif future_type == "albums":
results["albums"] = search_result["all"]
results["saved_albums"] = search_result["saved"]
user_ids.update(get_users_ids(search_result["all"]))
with db.scoped_session() as session:
# Add users back
users = get_users_by_id(session, list(user_ids), current_user_id)
for (_, result_list) in results.items():
for result in result_list:
user_id = None
if "playlist_owner_id" in result:
user_id = result["playlist_owner_id"]
elif "owner_id" in result:
user_id = result["owner_id"]
if user_id is not None:
user = users[user_id]
result["user"] = user
return extend_search(results)
