Skip to content

Commit

Permalink
fix: get granules pagination (#502)
Browse files Browse the repository at this point in the history 
* breaking: using latest uds-lib + update docker

* feat: use latest uds-lib

* fix: (in progress) redo granules search

* fix: pagination is working now. But total size is wrong (no solution)

* fix: adding back total count as there will be only 1 doc

* fix: sending real result size

* fix: wrong logic to get pagination

* fix: unsure where the error is

* fix: update page size bug

* fix: total size test case
  • Loading branch information
@wphyojpl
wphyojpl authored Jan 17, 2025
1 parent 70ff291 commit 9f5b5cf
Show file tree
Hide file tree
Showing 5 changed files with 54 additions and 14 deletions.
17 changes: 11 additions & 6 deletions cumulus_lambda_functions/lib/uds_db/granules_db_index.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -303,18 +303,21 @@ def add_entry(self, tenant: str, tenant_venue: str, json_body: dict, doc_id: str

def dsl_search(self, tenant: str, tenant_venue: str, search_dsl: dict):
read_alias_name = f'{DBConstants.granules_read_alias_prefix}_{tenant}_{tenant_venue}'.lower().strip()
if 'sort' not in search_dsl:
search_result = self.__es.query(search_dsl,
querying_index=read_alias_name) if 'sort' in search_dsl else self.__es.query(
search_dsl, querying_index=read_alias_name)
if 'sort' not in search_dsl: # We cannot paginate w/o sort. So, max is 10k items:
# This also assumes "size" should be part of search_dsl
search_result = self.__es.query(search_dsl, querying_index=read_alias_name)
LOGGER.debug(f'search_finished: {len(search_result["hits"]["hits"])}')
return search_result
# we can run paginate search
original_size = search_dsl['size'] if 'size' in search_dsl else 20
total_size = -999
result = []
duplicates = set([])
while len(result) < original_size:
search_dsl['size'] = (original_size - len(result)) * 2
search_result = self.__es.query_pages(search_dsl, querying_index=read_alias_name) if 'sort' in search_dsl else self.__es.query(search_dsl, querying_index=read_alias_name)
search_result = self.__es.query_pages(search_dsl, querying_index=read_alias_name)
if total_size == -999:
total_size = self.__es.get_result_size(search_result)
if len(search_result['hits']['hits']) < 1:
break
for each in search_result['hits']['hits']:
Expand All @@ -324,10 +327,12 @@ def dsl_search(self, tenant: str, tenant_venue: str, search_dsl: dict):
search_dsl['search_after'] = search_result['hits']['hits'][-1]['sort']

LOGGER.debug(f'search_finished: {len(result)}')
if len(result) > original_size:
result = result[:original_size]
return {
'hits': {
"total": {
"value": len(result)
"value": total_size,
},
'hits': result
}
Expand Down
6 changes: 3 additions & 3 deletions cumulus_lambda_functions/uds_api/dapa/granules_dapa_query_es.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ def __generate_es_dsl(self):
if self.__filter_input is not None:
query_terms.append(CqlParser('properties').transform(self.__filter_input))
query_dsl = {
'track_total_hits': True,
'track_total_hits': self.__offset is None,
'size': self.__limit,
# "collapse": {"field": "id"},
'sort': [
Expand Down Expand Up @@ -228,11 +228,11 @@ def start(self):
each_granules_query_result_stripped['links'].append(self_link)
self.__restructure_each_granule_result(each_granules_query_result_stripped)

pagination_link = '' if len(granules_query_result['hits']['hits']) < self.__limit else ','.join([k if isinstance(k, str) else str(k) for k in granules_query_result['hits']['hits'][-1]['sort']])
pagination_link = '' if len(granules_query_result['hits']['hits']) < 1 else ','.join([k if isinstance(k, str) else str(k) for k in granules_query_result['hits']['hits'][-1]['sort']])
return {
'statusCode': 200,
'body': {
'numberMatched': {'total_size': result_size},
'numberMatched': {'total_size': -1 if self.__offset is not None else result_size},
'numberReturned': len(granules_query_result['hits']['hits']),
'stac_version': '1.0.0',
'type': 'FeatureCollection', # TODO correct name?
Expand Down
4 changes: 2 additions & 2 deletions tests/cumulus_lambda_functions/lib/uds_db/test_granules_db_index.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ def test_01(self):
def test_02(self):
os.environ['ES_URL'] = 'https://vpc-uds-sbx-cumulus-es-qk73x5h47jwmela5nbwjte4yzq.us-west-2.es.amazonaws.com'
os.environ['ES_PORT'] = '9200'
self.tenant = 'UDS_MY_LOCAL_ARCHIVE_TEST' # 'uds_local_test' # 'uds_sandbox'
self.tenant = 'UDS_LOCAL_TEST_3' # 'uds_local_test' # 'uds_sandbox'
self.tenant_venue = 'DEV' # 'DEV1' # 'dev'
search_dsl = {
'track_total_hits': True,
Expand All @@ -48,7 +48,7 @@ def test_02(self):
],
'query': {
'bool': {
'must': {'term': {'collection': {'value': f'URN:NASA:UNITY:{self.tenant}:{self.tenant_venue}:UDS_UNIT_COLLECTION___2408290522'}}}
'must': {'term': {'collection': {'value': f'URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001'}}}
}
},
}
Expand Down
19 changes: 19 additions & 0 deletions tests/cumulus_lambda_functions/uds_api/dapa/test_granules_dapa_query_es.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,25 @@ def setUp(self) -> None:
self.collection_version = '001'
return

def test_start_01(self):
os.environ['ES_URL'] = 'vpc-uds-sbx-cumulus-es-qk73x5h47jwmela5nbwjte4yzq.us-west-2.es.amazonaws.com'
os.environ['ES_PORT'] = '9200'
self.tenant = 'UDS_LOCAL_TEST_3' # 'uds_local_test' # 'uds_sandbox'
self.tenant_venue = 'DEV' # 'DEV1' # 'dev'
self.collection_name = 'DDD-01' # 'uds_collection' # 'sbx_collection'
self.collection_version = '001'

collection_id = f'URN:NASA:UNITY:{self.tenant}:{self.tenant_venue}:{self.collection_name}___{self.collection_version}'
granule_id = f'{collection_id}:test_file09'

granules_dapa_query = GranulesDapaQueryEs(collection_id, 10, '1736291597733,URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file10', None, None, None, f'localhost/api-prefix')
granules_result = granules_dapa_query.start()
print(granules_result)
print([k['id'] for k in granules_result['body']['features']])
# ['URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file20', 'URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file19', 'URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file14', 'URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file17', 'URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file18', 'URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file12', 'URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file13', 'URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file15', 'URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file06', 'URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file01']
# ['URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file05', 'URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file03', 'URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file09', 'URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file16', 'URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file11', 'URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file04', 'URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file08', 'URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file02', 'URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file07', 'URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001:test_file10']
return

def test_get_single_granule_01(self):
os.environ['ES_URL'] = 'vpc-uds-sbx-cumulus-es-qk73x5h47jwmela5nbwjte4yzq.us-west-2.es.amazonaws.com'
os.environ['ES_PORT'] = '9200'
Expand Down
22 changes: 19 additions & 3 deletions tests/integration_tests/test_uds_api.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -184,8 +184,8 @@ def test_collections_get_single_granule(self):
return

def test_granules_get(self):
# post_url = f'{self.uds_url}collections/urn:nasa:unity:unity:dev:SBG-L2A_RFL___1/items/' # MCP Dev
post_url = f'{self.uds_url}collections/urn:nasa:unity:asips:int:P1590011-T___1/items/' # MCP OPS
post_url = f'{self.uds_url}collections/URN:NASA:UNITY:UDS_LOCAL_TEST_3:DEV:DDD-01___001/items/?limit=10' # MCP Dev
# post_url = f'{self.uds_url}collections/urn:nasa:unity:asips:int:P1590011-T___1/items/?limit=10' # MCP OPS
headers = {
'Authorization': f'Bearer {self.bearer_token}',
}
Expand All @@ -194,11 +194,27 @@ def test_granules_get(self):
headers=headers,
)
response_json = json.loads(query_result.text)
print(json.dumps(response_json, indent=4))
# print(json.dumps(response_json, indent=4))
print(f"length: {len(response_json['features'])}")
self.assertEqual(query_result.status_code, 200, f'wrong status code. {query_result.text}')
total_size = response_json['numberMatched']['total_size']
current_size = len(response_json['features'])
links = {k['rel']: k['href'] for k in response_json['links'] if k['rel'] != 'root'}
for k, v in links.items():
self.assertTrue(v.startswith(self.uds_url), f'missing stage: {self.stage} in {v} for {k}')

while len(response_json['features']) > 0:
get_next_link = [k['href'] for k in response_json['links'] if k['rel'] == 'next']
get_next_link = get_next_link[0]
print(get_next_link)
query_result = requests.get(url=get_next_link,
headers=headers,
)
response_json = json.loads(query_result.text)
self.assertEqual(query_result.status_code, 200, f'wrong status code. {query_result.text}')
current_size += len(response_json['features'])
print(f"length: {len(response_json['features'])}")
self.assertEqual(total_size, current_size, f'mismatched size')
return

def test_single_granule_get(self):
Expand Down

0 comments on commit 9f5b5cf

Please sign in to comment.