Skip to content

Commit 9edf047

Browse files
Yuri ZmytrakovYuri Zmytrakov
Yuri Zmytrakov
authored and
Yuri Zmytrakov
committed
temp
1 parent cb7ee7c commit 9edf047

File tree

4 files changed

+105
-120
lines changed

4 files changed

+105
-120
lines changed

‎stac_fastapi/core/stac_fastapi/core/core.py‎

Lines changed: 14 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -24,10 +24,10 @@
2424
from stac_fastapi.core.base_settings import ApiBaseSettings
2525
from stac_fastapi.core.datetime_utils import format_datetime_range
2626
from stac_fastapi.core.models.links import PagingLinks
27-
from stac_fastapi.core.redis_utils import connect_redis, get_prev_link, save_self_link
27+
from stac_fastapi.core.redis_utils import handle_pagination_links
2828
from stac_fastapi.core.serializers import CollectionSerializer, ItemSerializer
2929
from stac_fastapi.core.session import Session
30-
from stac_fastapi.core.utilities import filter_fields, get_bool_env
30+
from stac_fastapi.core.utilities import filter_fields
3131
from stac_fastapi.extensions.core.transaction import AsyncBaseTransactionsClient
3232
from stac_fastapi.extensions.core.transaction.request import (
3333
PartialCollection,
@@ -329,20 +329,6 @@ async def all_collections(
329329
if parsed_sort:
330330
sort = parsed_sort
331331

332-
current_url = str(request.url)
333-
redis_enable = get_bool_env("REDIS_ENABLE", default=False)
334-
335-
redis = None
336-
if redis_enable:
337-
try:
338-
redis = await connect_redis()
339-
logger.info("Redis connection established successfully")
340-
except Exception as e:
341-
redis = None
342-
logger.warning(
343-
f"Redis connection failed, continuing without Redis: {e}"
344-
)
345-
346332
# Convert q to a list if it's a string
347333
q_list = None
348334
if q is not None:
@@ -441,21 +427,12 @@ async def all_collections(
441427
},
442428
]
443429

444-
if redis_enable and redis:
445-
if next_token:
446-
await save_self_link(redis, next_token, current_url)
447-
448-
prev_link = await get_prev_link(redis, token)
449-
if prev_link:
450-
links.insert(
451-
0,
452-
{
453-
"rel": "prev",
454-
"type": "application/json",
455-
"method": "GET",
456-
"href": prev_link,
457-
},
458-
)
430+
await handle_pagination_links(
431+
current_url=str(request.url),
432+
token=token,
433+
next_token=next_token,
434+
links=links,
435+
)
459436

460437
if next_token:
461438
next_link = PagingLinks(next=next_token, request=request).link_next()
@@ -775,8 +752,6 @@ async def post_search(
775752
HTTPException: If there is an error with the cql2_json filter.
776753
"""
777754
base_url = str(request.base_url)
778-
redis_enable = get_bool_env("REDIS_ENABLE", default=False)
779-
780755
search = self.database.make_search()
781756

782757
if search_request.ids:
@@ -901,29 +876,12 @@ async def post_search(
901876
)
902877
links.extend(collection_links)
903878

904-
if redis_enable:
905-
redis = None
906-
try:
907-
redis = await connect_redis()
908-
logger.info("Redis connection established successfully")
909-
self_link = str(request.url)
910-
await save_self_link(redis, next_token, self_link)
911-
912-
prev_link = await get_prev_link(redis, token_param)
913-
if prev_link:
914-
links.insert(
915-
0,
916-
{
917-
"rel": "prev",
918-
"type": "application/json",
919-
"method": "GET",
920-
"href": prev_link,
921-
},
922-
)
923-
except Exception as e:
924-
logger.warning(
925-
f"Redis connection failed, continuing without Redis: {e}"
926-
)
879+
await handle_pagination_links(
880+
current_url=str(request.url),
881+
token=token_param,
882+
next_token=next_token,
883+
links=links,
884+
)
927885

928886
return stac_types.ItemCollection(
929887
type="FeatureCollection",

‎stac_fastapi/core/stac_fastapi/core/redis_utils.py‎

Lines changed: 91 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,18 @@
11
"""Utilities for connecting to and managing Redis connections."""
22

3+
import logging
34
from typing import Optional
45

56
from pydantic_settings import BaseSettings
67
from redis import asyncio as aioredis
78
from redis.asyncio.sentinel import Sentinel
89

10+
from stac_fastapi.core.utilities import get_bool_env
11+
912
redis_pool: Optional[aioredis.Redis] = None
1013

14+
logger = logging.getLogger(__name__)
15+
1116

1217
class RedisSentinelSettings(BaseSettings):
1318
"""Configuration for connecting to Redis Sentinel."""
@@ -25,7 +30,7 @@ class RedisSentinelSettings(BaseSettings):
2530

2631

2732
class RedisSettings(BaseSettings):
28-
"""Configuration for connecting Redis Sentinel."""
33+
"""Configuration for connecting Redis."""
2934

3035
REDIS_HOST: str = ""
3136
REDIS_PORT: int = 6379
@@ -38,75 +43,71 @@ class RedisSettings(BaseSettings):
3843
REDIS_HEALTH_CHECK_INTERVAL: int = 30
3944

4045

41-
# Select the Redis or Redis Sentinel configuration
42-
redis_settings: BaseSettings = RedisSettings()
43-
46+
# Configure only one Redis configuration
47+
sentinel_settings = RedisSentinelSettings()
48+
standalone_settings = RedisSettings()
4449

45-
async def connect_redis(settings: Optional[RedisSettings] = None) -> aioredis.Redis:
46-
"""Return a Redis connection."""
47-
global redis_pool
48-
settings = settings or redis_settings
49-
50-
if not settings.REDIS_HOST or not settings.REDIS_PORT:
51-
return None
5250

53-
if redis_pool is None:
54-
pool = aioredis.ConnectionPool(
55-
host=settings.REDIS_HOST,
56-
port=settings.REDIS_PORT,
57-
db=settings.REDIS_DB,
58-
max_connections=settings.REDIS_MAX_CONNECTIONS,
59-
decode_responses=settings.REDIS_DECODE_RESPONSES,
60-
retry_on_timeout=settings.REDIS_RETRY_TIMEOUT,
61-
health_check_interval=settings.REDIS_HEALTH_CHECK_INTERVAL,
62-
)
63-
redis_pool = aioredis.Redis(
64-
connection_pool=pool, client_name=settings.REDIS_CLIENT_NAME
65-
)
66-
return redis_pool
67-
68-
69-
async def connect_redis_sentinel(
70-
settings: Optional[RedisSentinelSettings] = None,
71-
) -> Optional[aioredis.Redis]:
72-
"""Return a Redis Sentinel connection."""
51+
async def connect_redis() -> Optional[aioredis.Redis]:
52+
"""Return a Redis connection Redis or Redis Sentinel."""
7353
global redis_pool
7454

75-
settings = settings or redis_settings
55+
if redis_pool is not None:
56+
return redis_pool
57+
58+
try:
59+
if sentinel_settings.REDIS_SENTINEL_HOSTS:
60+
hosts = [
61+
h.strip()
62+
for h in sentinel_settings.REDIS_SENTINEL_HOSTS.split(",")
63+
if h.strip()
64+
]
65+
ports = [
66+
int(p.strip())
67+
for p in sentinel_settings.REDIS_SENTINEL_PORTS.split(",")
68+
if p.strip()
69+
]
7670

77-
if (
78-
not settings.REDIS_SENTINEL_HOSTS
79-
or not settings.REDIS_SENTINEL_PORTS
80-
or not settings.REDIS_SENTINEL_MASTER_NAME
81-
):
82-
return None
83-
84-
hosts = [h.strip() for h in settings.REDIS_SENTINEL_HOSTS.split(",") if h.strip()]
85-
ports = [
86-
int(p.strip()) for p in settings.REDIS_SENTINEL_PORTS.split(",") if p.strip()
87-
]
88-
89-
if redis_pool is None:
90-
try:
9171
sentinel = Sentinel(
9272
[(h, p) for h, p in zip(hosts, ports)],
93-
decode_responses=settings.REDIS_DECODE_RESPONSES,
94-
)
95-
master = sentinel.master_for(
96-
service_name=settings.REDIS_SENTINEL_MASTER_NAME,
97-
db=settings.REDIS_DB,
98-
decode_responses=settings.REDIS_DECODE_RESPONSES,
99-
retry_on_timeout=settings.REDIS_RETRY_TIMEOUT,
100-
client_name=settings.REDIS_CLIENT_NAME,
101-
max_connections=settings.REDIS_MAX_CONNECTIONS,
102-
health_check_interval=settings.REDIS_HEALTH_CHECK_INTERVAL,
73+
decode_responses=sentinel_settings.REDIS_DECODE_RESPONSES,
10374
)
104-
redis_pool = master
10575

106-
except Exception:
76+
redis_pool = sentinel.master_for(
77+
service_name=sentinel_settings.REDIS_SENTINEL_MASTER_NAME,
78+
db=sentinel_settings.REDIS_DB,
79+
decode_responses=sentinel_settings.REDIS_DECODE_RESPONSES,
80+
retry_on_timeout=sentinel_settings.REDIS_RETRY_TIMEOUT,
81+
client_name=sentinel_settings.REDIS_CLIENT_NAME,
82+
max_connections=sentinel_settings.REDIS_MAX_CONNECTIONS,
83+
health_check_interval=sentinel_settings.REDIS_HEALTH_CHECK_INTERVAL,
84+
)
85+
logger.info("Connected to Redis Sentinel")
86+
87+
elif standalone_settings.REDIS_HOST:
88+
pool = aioredis.ConnectionPool(
89+
host=standalone_settings.REDIS_HOST,
90+
port=standalone_settings.REDIS_PORT,
91+
db=standalone_settings.REDIS_DB,
92+
max_connections=standalone_settings.REDIS_MAX_CONNECTIONS,
93+
decode_responses=standalone_settings.REDIS_DECODE_RESPONSES,
94+
retry_on_timeout=standalone_settings.REDIS_RETRY_TIMEOUT,
95+
health_check_interval=standalone_settings.REDIS_HEALTH_CHECK_INTERVAL,
96+
)
97+
redis_pool = aioredis.Redis(
98+
connection_pool=pool, client_name=standalone_settings.REDIS_CLIENT_NAME
99+
)
100+
logger.info("Connected to Redis")
101+
else:
102+
logger.warning("No Redis configuration found")
107103
return None
108104

109-
return redis_pool
105+
return redis_pool
106+
107+
except Exception as e:
108+
logger.error(f"Failed to connect to Redis: {e}")
109+
redis_pool = None
110+
return None
110111

111112

112113
async def save_self_link(
@@ -122,3 +123,34 @@ async def get_prev_link(redis: aioredis.Redis, token: Optional[str]) -> Optional
122123
if not token:
123124
return None
124125
return await redis.get(f"nav:self:{token}")
126+
127+
128+
async def handle_pagination_links(
129+
current_url: str, token: str, next_token: str, links: list
130+
) -> None:
131+
"""Handle Redis pagination."""
132+
redis_enable = get_bool_env("REDIS_ENABLE", default=False)
133+
redis = None
134+
if redis_enable:
135+
try:
136+
redis = await connect_redis()
137+
logger.info("Redis connection established successfully")
138+
except Exception as e:
139+
redis = None
140+
logger.warning(f"Redis connection failed: {e}")
141+
142+
if redis_enable and redis:
143+
if next_token:
144+
await save_self_link(redis, next_token, current_url)
145+
146+
prev_link = await get_prev_link(redis, token)
147+
if prev_link:
148+
links.insert(
149+
0,
150+
{
151+
"rel": "prev",
152+
"type": "application/json",
153+
"method": "GET",
154+
"href": prev_link,
155+
},
156+
)

‎stac_fastapi/tests/redis/test_redis_pagination.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,6 @@ async def test_search_pagination_uses_redis_cache(
3232
)
3333
next_token = next_link["body"]["token"]
3434

35-
# Expect the previous link on the second page to be retrieved from Redis cache
3635
resp2 = await app_client.post(
3736
"/search",
3837
json={"collections": [collection_id], "limit": 1, "token": next_token},
@@ -67,7 +66,6 @@ async def test_collections_pagination_uses_redis_cache(
6766
)
6867
next_token = next_link["href"].split("token=")[1]
6968

70-
# Expect the previous link on the second page to be retrieved from Redis cache
7169
resp2 = await app_client.get(
7270
"/collections", params={"limit": 1, "token": next_token}
7371
)

‎stac_fastapi/tests/redis/test_redis_utils.py‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,16 +9,13 @@ async def test_redis_connection():
99
redis = await connect_redis()
1010
assert redis is not None
1111

12-
# Test set/get
1312
await redis.set("string_key", "string_value")
1413
string_value = await redis.get("string_key")
1514
assert string_value == "string_value"
1615

17-
# Test key retrieval operation
1816
exists = await redis.exists("string_key")
1917
assert exists == 1
2018

21-
# Test key deletion
2219
await redis.delete("string_key")
2320
deleted_value = await redis.get("string_key")
2421
assert deleted_value is None

0 commit comments

Comments
 (0)