Skip to content

Commit 1a78326

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

File tree

4 files changed

+105
-121
lines changed

4 files changed

+105
-121
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 & 60 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
@@ -37,76 +42,71 @@ class RedisSettings(BaseSettings):
3742
REDIS_CLIENT_NAME: str = "stac-fastapi-app"
3843
REDIS_HEALTH_CHECK_INTERVAL: int = 30
3944

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

4449

45-
async def connect_redis(settings: Optional[RedisSettings] = None) -> aioredis.Redis:
46-
"""Return a Redis connection."""
50+
async def connect_redis() -> Optional[aioredis.Redis]:
51+
"""Return a Redis connection Redis or Redis Sentinel."""
4752
global redis_pool
48-
settings = settings or redis_settings
4953

50-
if not settings.REDIS_HOST or not settings.REDIS_PORT:
51-
return None
52-
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."""
73-
global redis_pool
54+
if redis_pool is not None:
55+
return redis_pool
56+
57+
try:
58+
if sentinel_settings.REDIS_SENTINEL_HOSTS:
59+
hosts = [
60+
h.strip()
61+
for h in sentinel_settings.REDIS_SENTINEL_HOSTS.split(",")
62+
if h.strip()
63+
]
64+
ports = [
65+
int(p.strip())
66+
for p in sentinel_settings.REDIS_SENTINEL_PORTS.split(",")
67+
if p.strip()
68+
]
7469

75-
settings = settings or redis_settings
76-
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:
9170
sentinel = Sentinel(
9271
[(h, p) for h, p in zip(hosts, ports)],
93-
decode_responses=settings.REDIS_DECODE_RESPONSES,
72+
decode_responses=sentinel_settings.REDIS_DECODE_RESPONSES,
9473
)
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,
103-
)
104-
redis_pool = master
10574

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

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

111111

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

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