Skip to content

Commit 7b896f7

Browse files
Yuri ZmytrakovYuri Zmytrakov
Yuri Zmytrakov
authored and
Yuri Zmytrakov
committed
temp
1 parent cb7ee7c commit 7b896f7

File tree

4 files changed

+108
-118
lines changed

4 files changed

+108
-118
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: 94 additions & 57 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."""
@@ -38,75 +43,76 @@ 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+
# Create settings instances
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, automatically detecting Sentinel or standalone."""
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+
# Check if Sentinel is configured
60+
if sentinel_settings.REDIS_SENTINEL_HOSTS:
61+
# Use Sentinel connection
62+
hosts = [
63+
h.strip()
64+
for h in sentinel_settings.REDIS_SENTINEL_HOSTS.split(",")
65+
if h.strip()
66+
]
67+
ports = [
68+
int(p.strip())
69+
for p in sentinel_settings.REDIS_SENTINEL_PORTS.split(",")
70+
if p.strip()
71+
]
7672

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:
9173
sentinel = Sentinel(
9274
[(h, p) for h, p in zip(hosts, ports)],
93-
decode_responses=settings.REDIS_DECODE_RESPONSES,
75+
decode_responses=sentinel_settings.REDIS_DECODE_RESPONSES,
9476
)
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,
77+
78+
redis_pool = sentinel.master_for(
79+
service_name=sentinel_settings.REDIS_SENTINEL_MASTER_NAME,
80+
db=sentinel_settings.REDIS_DB,
81+
decode_responses=sentinel_settings.REDIS_DECODE_RESPONSES,
82+
retry_on_timeout=sentinel_settings.REDIS_RETRY_TIMEOUT,
83+
client_name=sentinel_settings.REDIS_CLIENT_NAME,
84+
max_connections=sentinel_settings.REDIS_MAX_CONNECTIONS,
85+
health_check_interval=sentinel_settings.REDIS_HEALTH_CHECK_INTERVAL,
10386
)
104-
redis_pool = master
10587

106-
except Exception:
88+
logger.info("Connected to Redis Sentinel")
89+
90+
elif standalone_settings.REDIS_HOST:
91+
# Use standalone Redis connection
92+
pool = aioredis.ConnectionPool(
93+
host=standalone_settings.REDIS_HOST,
94+
port=standalone_settings.REDIS_PORT,
95+
db=standalone_settings.REDIS_DB,
96+
max_connections=standalone_settings.REDIS_MAX_CONNECTIONS,
97+
decode_responses=standalone_settings.REDIS_DECODE_RESPONSES,
98+
retry_on_timeout=standalone_settings.REDIS_RETRY_TIMEOUT,
99+
health_check_interval=standalone_settings.REDIS_HEALTH_CHECK_INTERVAL,
100+
)
101+
redis_pool = aioredis.Redis(
102+
connection_pool=pool, client_name=standalone_settings.REDIS_CLIENT_NAME
103+
)
104+
105+
logger.info("Connected to standalone Redis")
106+
else:
107+
logger.warning("No Redis configuration found")
107108
return None
108109

109-
return redis_pool
110+
return redis_pool
111+
112+
except Exception as e:
113+
logger.error(f"Failed to connect to Redis: {e}")
114+
redis_pool = None
115+
return None
110116

111117

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

‎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)