From 2897708e7204a2a3fe1abbf64f6881c95613877c Mon Sep 17 00:00:00 2001 From: CahidArda Date: Wed, 21 Jan 2026 12:00:16 +0300 Subject: [PATCH 01/12] feat: add redis search --- tests/commands/search/test_search_advanced.py | 379 +++++++++++++ tests/commands/search/test_search_index.py | 497 ++++++++++++++++++ tests/commands/search/test_search_utils.py | 390 ++++++++++++++ upstash_redis/asyncio/client.py | 6 + upstash_redis/client.py | 6 + upstash_redis/search_index.py | 281 ++++++++++ upstash_redis/search_namespace.py | 113 ++++ upstash_redis/search_types.py | 158 ++++++ upstash_redis/search_utils.py | 440 ++++++++++++++++ 9 files changed, 2270 insertions(+) create mode 100644 tests/commands/search/test_search_advanced.py create mode 100644 tests/commands/search/test_search_index.py create mode 100644 tests/commands/search/test_search_utils.py create mode 100644 upstash_redis/search_index.py create mode 100644 upstash_redis/search_namespace.py create mode 100644 upstash_redis/search_types.py create mode 100644 upstash_redis/search_utils.py diff --git a/tests/commands/search/test_search_advanced.py b/tests/commands/search/test_search_advanced.py new file mode 100644 index 0000000..6ad8cea --- /dev/null +++ b/tests/commands/search/test_search_advanced.py @@ -0,0 +1,379 @@ +"""Tests for JSON and nested schema search functionality.""" + +import json +from typing import List, TypedDict, Generator + +import pytest + +from upstash_redis import Redis +from upstash_redis.search_index import SearchIndex + + +class IndexFixture(TypedDict): + """Type definition for index fixtures.""" + index: SearchIndex + redis: Redis + keys: List[str] + name: str + + +def random_id() -> str: + """Generate a random ID for test isolation.""" + import random + import string + + return "".join(random.choices(string.ascii_lowercase + string.digits, k=8)) + + +@pytest.fixture +def redis_client(): + """Create a Redis client for testing.""" + return Redis.from_env() + + +class TestQueryJson: + """Tests for querying JSON indexes.""" + + @pytest.fixture(scope="class") + def json_index(self) -> Generator[IndexFixture, None, None]: + """Create a JSON index with test data.""" + name = f"test-query-json-{random_id()}" + prefix = f"{name}:" + keys = [] + + redis = Redis.from_env() + + schema = { + "name": "TEXT", + "description": "TEXT", + "category": {"type": "TEXT", "noTokenize": True}, + "price": {"type": "F64", "fast": True}, + "stock": {"type": "U64", "fast": True}, + "active": "BOOL", + } + + index = redis.search.createIndex( + name=name, schema=schema, dataType="json", prefix=prefix + ) + + # Add test data + test_data = [ + { + "name": "Laptop Pro", + "description": "High performance laptop", + "category": "electronics", + "price": 1299.99, + "stock": 50, + "active": True, + }, + { + "name": "Laptop Basic", + "description": "Budget friendly laptop", + "category": "electronics", + "price": 599.99, + "stock": 100, + "active": True, + }, + { + "name": "Wireless Mouse", + "description": "Ergonomic wireless mouse", + "category": "electronics", + "price": 29.99, + "stock": 200, + "active": True, + }, + { + "name": "USB Cable", + "description": "Fast charging USB cable", + "category": "accessories", + "price": 9.99, + "stock": 500, + "active": True, + }, + ] + + for i, datum in enumerate(test_data): + key = f"{prefix}{i}" + keys.append(key) + redis.json.set(key, "$", datum) + + index.wait_indexing() + + yield {"index": index, "redis": redis, "keys": keys, "name": name} + + # Cleanup + try: + index.drop() + except: + pass + if keys: + redis.delete(*keys) + + def test_query_json_text_eq(self, json_index: IndexFixture): + """Test querying JSON index with $eq on text field.""" + index = json_index["index"] + result = index.query(filter={"name": {"$eq": "Laptop"}}) + + assert len(result) > 0 + + def test_query_json_fuzzy(self, json_index: IndexFixture): + """Test querying JSON index with fuzzy search.""" + index = json_index["index"] + result = index.query(filter={"name": {"$fuzzy": "laptopp"}}) + + assert len(result) > 0 + + def test_query_json_phrase(self, json_index: IndexFixture): + """Test querying JSON index with phrase matching.""" + index = json_index["index"] + result = index.query(filter={"description": {"$phrase": "wireless mouse"}}) + + assert len(result) > 0 + + def test_query_json_regex(self, json_index: IndexFixture): + """Test querying JSON index with regex pattern.""" + index = json_index["index"] + result = index.query(filter={"name": {"$regex": "Laptop.*"}}) + + assert len(result) > 0 + + def test_query_json_numeric_gt(self, json_index: IndexFixture): + """Test querying JSON index with $gt on numeric field.""" + index = json_index["index"] + result = index.query(filter={"price": {"$gt": 500}}) + + assert len(result) > 0 + + def test_query_json_with_sorting(self, json_index: IndexFixture): + """Test querying JSON index with sorting.""" + index = json_index["index"] + result = index.query( + filter={"category": {"$eq": "electronics"}}, + orderBy={"price": "DESC"}, + limit=3, + ) + + assert len(result) > 0 + + +class TestHashIndex: + """Tests for hash index queries.""" + + @pytest.fixture(scope="class") + def hash_index(self) -> Generator[IndexFixture, None, None]: + """Create a hash index with test data.""" + name = f"test-hash-query-{random_id()}" + prefix = f"{name}:" + keys = [] + + redis = Redis.from_env() + + schema = {"name": "TEXT", "score": {"type": "U64", "fast": True}} + + index = redis.search.createIndex( + name=name, schema=schema, dataType="hash", prefix=prefix + ) + + # Add test data using HSET + test_data = [ + {"name": "Alice", "score": "95"}, + {"name": "Bob", "score": "87"}, + {"name": "Charlie", "score": "92"}, + ] + + for i, datum in enumerate(test_data): + key = f"{prefix}{i}" + keys.append(key) + redis.hset(key, values=datum) + + index.wait_indexing() + + yield {"index": index, "redis": redis, "keys": keys, "name": name} + + # Cleanup + try: + index.drop() + except: + pass + if keys: + redis.delete(*keys) + + def test_query_hash_by_text(self, hash_index: IndexFixture): + """Test querying hash index by text field.""" + index = hash_index["index"] + result = index.query(filter={"name": {"$eq": "Alice"}}) + + assert len(result) > 0 + + def test_query_hash_with_sorting(self, hash_index: IndexFixture): + """Test querying hash index with sorting.""" + index = hash_index["index"] + result = index.query(filter={"score": {"$gte": 80}}, orderBy={"score": "DESC"}) + + assert len(result) > 0 + + +class TestNestedStringIndex: + """Tests for nested string index queries.""" + + @pytest.fixture(scope="class") + def nested_string_index(self) -> Generator[IndexFixture, None, None]: + """Create a nested string index with test data.""" + name = f"test-nested-string-{random_id()}" + prefix = f"{name}:" + keys = [] + + redis = Redis.from_env() + + schema = { + "title": "TEXT", + "author": {"name": "TEXT", "email": "TEXT"}, + "stats": {"views": {"type": "U64", "fast": True}, "likes": {"type": "U64", "fast": True}}, + } + + index = redis.search.createIndex( + name=name, schema=schema, dataType="string", prefix=prefix + ) + + test_data = [ + { + "title": "First Post", + "author": {"name": "John Doe", "email": "john@example.com"}, + "stats": {"views": 1000, "likes": 50}, + }, + { + "title": "Second Post", + "author": {"name": "Jane Smith", "email": "jane@example.com"}, + "stats": {"views": 500, "likes": 30}, + }, + { + "title": "Third Post", + "author": {"name": "John Doe", "email": "john@example.com"}, + "stats": {"views": 2000, "likes": 100}, + }, + ] + + for i, datum in enumerate(test_data): + key = f"{prefix}{i}" + keys.append(key) + redis.set(key, json.dumps(datum)) + + index.wait_indexing() + + yield {"index": index, "redis": redis, "keys": keys, "name": name} + + # Cleanup + try: + index.drop() + except: + pass + if keys: + redis.delete(*keys) + + def test_query_nested_text_field(self, nested_string_index: IndexFixture): + """Test querying nested text field.""" + index = nested_string_index["index"] + result = index.query(filter={"author.name": {"$eq": "John"}}) + + assert len(result) > 0 + + def test_query_nested_numeric_field(self, nested_string_index: IndexFixture): + """Test querying nested numeric field.""" + index = nested_string_index["index"] + result = index.query(filter={"stats.views": {"$eq": 1000}}) + + assert len(result) > 0 + + def test_query_nested_with_sorting(self, nested_string_index: IndexFixture): + """Test querying with sorting on nested field.""" + index = nested_string_index["index"] + result = index.query( + filter={"author.name": {"$eq": "John"}}, + select={"author.email": True}, + orderBy={"stats.views": "DESC"}, + ) + + assert len(result) > 0 + + +class TestNestedJsonIndex: + """Tests for nested JSON index queries.""" + + @pytest.fixture(scope="class") + def nested_json_index(self) -> Generator[IndexFixture, None, None]: + """Create a nested JSON index with test data.""" + name = f"test-nested-json-{random_id()}" + prefix = f"{name}:" + keys = [] + + redis = Redis.from_env() + + schema = { + "title": "TEXT", + "author": {"name": "TEXT", "email": "TEXT"}, + "stats": {"views": {"type": "U64", "fast": True}, "likes": {"type": "U64", "fast": True}}, + } + + index = redis.search.createIndex( + name=name, schema=schema, dataType="json", prefix=prefix + ) + + test_data = [ + { + "title": "First Post", + "author": {"name": "John Doe", "email": "john@example.com"}, + "stats": {"views": 1000, "likes": 50}, + }, + { + "title": "Second Post", + "author": {"name": "Jane Smith", "email": "jane@example.com"}, + "stats": {"views": 500, "likes": 30}, + }, + { + "title": "Third Post", + "author": {"name": "John Doe", "email": "john@example.com"}, + "stats": {"views": 2000, "likes": 100}, + }, + ] + + for i, datum in enumerate(test_data): + key = f"{prefix}{i}" + keys.append(key) + redis.json.set(key, "$", datum) + + index.wait_indexing() + + yield {"index": index, "redis": redis, "keys": keys, "name": name} + + # Cleanup + try: + index.drop() + except: + pass + if keys: + redis.delete(*keys) + + def test_query_nested_json_text_field(self, nested_json_index: IndexFixture): + """Test querying nested text field in JSON index.""" + index = nested_json_index["index"] + result = index.query(filter={"author.name": {"$eq": "John"}}) + + assert len(result) > 0 + + def test_query_nested_json_numeric_field(self, nested_json_index: IndexFixture): + """Test querying nested numeric field in JSON index.""" + index = nested_json_index["index"] + result = index.query(filter={"stats.views": {"$eq": 1000}}) + + assert len(result) > 0 + + def test_query_nested_json_with_sorting(self, nested_json_index: IndexFixture): + """Test querying with sorting on nested field in JSON index.""" + index = nested_json_index["index"] + result = index.query( + filter={"author.name": {"$eq": "John"}}, + select={"author.email": True}, + orderBy={"stats.views": "DESC"}, + ) + + assert len(result) > 0 diff --git a/tests/commands/search/test_search_index.py b/tests/commands/search/test_search_index.py new file mode 100644 index 0000000..29ca883 --- /dev/null +++ b/tests/commands/search/test_search_index.py @@ -0,0 +1,497 @@ +"""Tests for search index functionality.""" + +import json +import time +from typing import Dict, Any, List, TypedDict, Generator + +import pytest + +from upstash_redis import Redis +from upstash_redis.search_index import SearchIndex + + +class StringIndexFixture(TypedDict): + """Type definition for string_index fixture.""" + index: SearchIndex + redis: Redis + keys: List[str] + name: str + + +class CountIndexFixture(TypedDict): + """Type definition for count_index fixture.""" + index: SearchIndex + redis: Redis + keys: List[str] + name: str + + +def random_id() -> str: + """Generate a random ID for test isolation.""" + import random + import string + + return "".join(random.choices(string.ascii_lowercase + string.digits, k=8)) + + +@pytest.fixture +def redis_client(): + """Create a Redis client for testing.""" + return Redis.from_env() + + +@pytest.fixture +def cleanup_indexes(): + """Track and cleanup indexes created during tests.""" + indexes = [] + + yield indexes + + # Cleanup after test + redis = Redis.from_env() + for index_name in indexes: + try: + index = redis.search.index(index_name) + index.drop() + except: + pass + + +class TestCreateIndex: + """Tests for creating search indexes.""" + + def test_create_string_index_simple_schema(self, redis_client, cleanup_indexes): + """Test creating a string index with simple schema.""" + name = f"test-string-{random_id()}" + cleanup_indexes.append(name) + + schema = {"name": "TEXT", "age": {"type": "U64", "fast": True}} + + index = redis_client.search.createIndex( + name=name, schema=schema, dataType="string", prefix=f"{name}:" + ) + + assert index is not None + assert index.name == name + + def test_create_string_index_nested_schema(self, redis_client, cleanup_indexes): + """Test creating a string index with nested schema.""" + name = f"test-string-nested-{random_id()}" + cleanup_indexes.append(name) + + schema = { + "title": "TEXT", + "metadata": {"author": "TEXT", "views": {"type": "U64", "fast": True}}, + } + + index = redis_client.search.createIndex( + name=name, schema=schema, dataType="string", prefix=f"{name}:" + ) + + assert index is not None + + def test_create_hash_index(self, redis_client, cleanup_indexes): + """Test creating a hash index.""" + name = f"test-hash-{random_id()}" + cleanup_indexes.append(name) + + schema = {"title": "TEXT", "count": {"type": "U64", "fast": True}} + + index = redis_client.search.createIndex( + name=name, schema=schema, dataType="hash", prefix=f"{name}:" + ) + + assert index is not None + + def test_create_json_index_simple_schema(self, redis_client, cleanup_indexes): + """Test creating a JSON index with simple schema.""" + name = f"test-json-{random_id()}" + cleanup_indexes.append(name) + + schema = {"name": "TEXT", "age": {"type": "U64", "fast": True}} + + index = redis_client.search.createIndex( + name=name, schema=schema, dataType="json", prefix=f"{name}:" + ) + + assert index is not None + assert index.name == name + + def test_create_json_index_nested_schema(self, redis_client, cleanup_indexes): + """Test creating a JSON index with nested schema.""" + name = f"test-json-nested-{random_id()}" + cleanup_indexes.append(name) + + schema = { + "title": "TEXT", + "metadata": {"author": "TEXT", "views": {"type": "U64", "fast": True}}, + } + + index = redis_client.search.createIndex( + name=name, schema=schema, dataType="json", prefix=f"{name}:" + ) + + assert index is not None + + def test_create_index_with_language(self, redis_client, cleanup_indexes): + """Test creating an index with language option.""" + name = f"test-lang-{random_id()}" + cleanup_indexes.append(name) + + schema = {"content": "TEXT"} + + index = redis_client.search.createIndex( + name=name, + schema=schema, + dataType="string", + prefix=f"{name}:", + language="turkish", + ) + + assert index is not None + + def test_create_index_with_multiple_prefixes(self, redis_client, cleanup_indexes): + """Test creating an index with multiple prefixes.""" + name = f"test-multi-prefix-{random_id()}" + cleanup_indexes.append(name) + + schema = {"name": "TEXT"} + + index = redis_client.search.createIndex( + name=name, + schema=schema, + dataType="hash", + prefix=[f"{name}:users:", f"{name}:profiles:"], + ) + + assert index is not None + + +class TestQueryString: + """Tests for querying string indexes.""" + + @pytest.fixture(scope="class") + def string_index(self) -> Generator[StringIndexFixture, None, None]: + """Create a string index with test data.""" + name = f"test-query-string-{random_id()}" + prefix = f"{name}:" + keys = [] + + redis = Redis.from_env() + + schema = { + "name": "TEXT", + "description": "TEXT", + "category": {"type": "TEXT", "noTokenize": True}, + "price": {"type": "F64", "fast": True}, + "stock": {"type": "U64", "fast": True}, + "active": "BOOL", + } + + index = redis.search.createIndex( + name=name, schema=schema, dataType="string", prefix=prefix + ) + + # Add test data + test_data = [ + { + "name": "Laptop Pro", + "description": "High performance laptop", + "category": "electronics", + "price": 1299.99, + "stock": 50, + "active": True, + }, + { + "name": "Laptop Basic", + "description": "Budget friendly laptop", + "category": "electronics", + "price": 599.99, + "stock": 100, + "active": True, + }, + { + "name": "Wireless Mouse", + "description": "Ergonomic wireless mouse", + "category": "electronics", + "price": 29.99, + "stock": 200, + "active": True, + }, + { + "name": "USB Cable", + "description": "Fast charging USB cable", + "category": "accessories", + "price": 9.99, + "stock": 500, + "active": True, + }, + { + "name": "Phone Case", + "description": "Protective phone case", + "category": "accessories", + "price": 19.99, + "stock": 300, + "active": False, + }, + ] + + for i, datum in enumerate(test_data): + key = f"{prefix}{i}" + keys.append(key) + redis.set(key, json.dumps(datum)) + + # Wait for indexing + index.wait_indexing() + + yield {"index": index, "redis": redis, "keys": keys, "name": name} + + # Cleanup + try: + index.drop() + except: + pass + if keys: + redis.delete(*keys) + + def test_query_text_eq(self, string_index: StringIndexFixture): + """Test querying with $eq on text field.""" + index = string_index["index"] + result = index.query(filter={"name": {"$eq": "Laptop"}}) + + assert len(result) > 0 + + def test_query_text_fuzzy(self, string_index: StringIndexFixture): + """Test querying with fuzzy search for typo tolerance.""" + index = string_index["index"] + result = index.query(filter={"name": {"$fuzzy": "laptopp"}}) + + assert len(result) > 0 + + def test_query_text_phrase(self, string_index: StringIndexFixture): + """Test querying with phrase matching.""" + index = string_index["index"] + result = index.query(filter={"description": {"$phrase": "wireless mouse"}}) + + assert len(result) > 0 + + def test_query_text_regex(self, string_index: StringIndexFixture): + """Test querying with regex pattern.""" + index = string_index["index"] + result = index.query(filter={"name": {"$regex": "Laptop.*"}}) + + assert len(result) > 0 + + def test_query_numeric_gt(self, string_index: StringIndexFixture): + """Test querying with $gt on numeric field.""" + index = string_index["index"] + result = index.query(filter={"price": {"$gt": 500}}) + + assert len(result) > 0 + + def test_query_numeric_gte(self, string_index: StringIndexFixture): + """Test querying with $gte on numeric field.""" + index = string_index["index"] + result = index.query(filter={"stock": {"$gte": 100}}) + + assert len(result) > 0 + + def test_query_numeric_lt(self, string_index: StringIndexFixture): + """Test querying with $lt on numeric field.""" + index = string_index["index"] + result = index.query(filter={"price": {"$lt": 50}}) + + assert len(result) > 0 + + def test_query_numeric_lte(self, string_index: StringIndexFixture): + """Test querying with $lte on numeric field.""" + index = string_index["index"] + result = index.query(filter={"stock": {"$lte": 50}}) + + assert len(result) > 0 + + def test_query_boolean_eq(self, string_index: StringIndexFixture): + """Test querying with $eq on boolean field.""" + index = string_index["index"] + result = index.query(filter={"active": {"$eq": False}}) + + assert len(result) > 0 + + def test_query_with_limit(self, string_index: StringIndexFixture): + """Test querying with limit option.""" + index = string_index["index"] + result = index.query(filter={"category": {"$eq": "electronics"}}, limit=2) + + assert len(result) > 0 + assert len(result) <= 2 + + def test_query_with_pagination(self, string_index: StringIndexFixture): + """Test querying with offset for pagination.""" + index = string_index["index"] + first_page = index.query( + filter={"category": {"$eq": "electronics"}}, limit=2, offset=0 + ) + + second_page = index.query( + filter={"category": {"$eq": "electronics"}}, limit=2, offset=2 + ) + + assert len(first_page) > 0 + # second_page might be empty if there are only 2-3 results + + def test_query_with_sort_asc(self, string_index: StringIndexFixture): + """Test querying with sortBy ascending.""" + index = string_index["index"] + result = index.query( + filter={"category": {"$eq": "electronics"}}, + orderBy={"price": "ASC"}, + limit=3, + ) + + assert len(result) > 0 + + def test_query_with_sort_desc(self, string_index: StringIndexFixture): + """Test querying with sortBy descending.""" + index = string_index["index"] + result = index.query( + filter={"category": {"$eq": "electronics"}}, + orderBy={"price": "DESC"}, + limit=3, + ) + + assert len(result) > 0 + + def test_query_no_content(self, string_index: StringIndexFixture): + """Test querying with noContent (keys only).""" + index = string_index["index"] + result = index.query(filter={"category": {"$eq": "electronics"}}, select={}) + + assert len(result) > 0 + + def test_query_with_return_fields(self, string_index: StringIndexFixture): + """Test querying with specific return fields.""" + index = string_index["index"] + result = index.query( + filter={"category": {"$eq": "electronics"}}, + select={"category": True}, + highlight={"fields": []}, + ) + + assert len(result) > 0 + + +class TestCount: + """Tests for count functionality.""" + + @pytest.fixture(scope="class") + def count_index(self) -> Generator[CountIndexFixture, None, None]: + """Create an index for count testing.""" + name = f"test-count-{random_id()}" + prefix = f"{name}:" + keys = [] + + redis = Redis.from_env() + + schema = {"type": "TEXT", "value": {"type": "U64", "fast": True}} + + index = redis.search.createIndex( + name=name, schema=schema, dataType="string", prefix=prefix + ) + + # Add test data + for i in range(10): + key = f"{prefix}{i}" + keys.append(key) + redis.set(key, json.dumps({"type": "A" if i < 5 else "B", "value": i})) + + index.wait_indexing() + + yield {"index": index, "redis": redis, "keys": keys, "name": name} + + # Cleanup + try: + index.drop() + except: + pass + if keys: + redis.delete(*keys) + + def test_count_matching_documents(self, count_index: CountIndexFixture): + """Test counting matching documents.""" + index = count_index["index"] + result = index.count({"type": {"$eq": "A"}}) + + assert "count" in result + assert result["count"] > 0 + + def test_count_with_numeric_filter(self, count_index: CountIndexFixture): + """Test counting with numeric filter.""" + index = count_index["index"] + result = index.count({"value": {"$eq": 5}}) + + assert "count" in result + assert result["count"] > 0 + + +class TestDescribe: + """Tests for describe functionality.""" + + def test_describe_index(self, redis_client, cleanup_indexes): + """Test describing an index structure.""" + name = f"test-describe-{random_id()}" + cleanup_indexes.append(name) + + schema = { + "title": {"type": "TEXT", "noStem": True}, + "count": {"type": "U64", "fast": True}, + "active": "BOOL", + } + + index = redis_client.search.createIndex( + name=name, schema=schema, dataType="string", prefix=f"{name}:" + ) + + description = index.describe() + + assert description is not None + assert "name" in description or "schema" in description + + +class TestDrop: + """Tests for drop functionality.""" + + def test_drop_existing_index(self, redis_client): + """Test dropping an existing index.""" + name = f"test-drop-{random_id()}" + + schema = {"name": "TEXT"} + + index = redis_client.search.createIndex( + name=name, schema=schema, dataType="string", prefix=f"{name}:" + ) + + result = index.drop() + # Result should be 1 or 0 or "OK" + assert result is not None + + +class TestIndexMethod: + """Tests for redis.search.index method.""" + + def test_index_without_schema(self, redis_client): + """Test creating a SearchIndex instance without schema.""" + index = redis_client.search.index("test-index") + + assert index is not None + assert index.name == "test-index" + assert index.schema is None + + def test_index_with_schema(self, redis_client): + """Test creating a SearchIndex instance with schema.""" + schema = {"name": "TEXT", "age": {"type": "U64", "fast": True}} + + index = redis_client.search.index("test-index", schema) + + assert index is not None + assert index.name == "test-index" + assert index.schema == schema diff --git a/tests/commands/search/test_search_utils.py b/tests/commands/search/test_search_utils.py new file mode 100644 index 0000000..ec028ec --- /dev/null +++ b/tests/commands/search/test_search_utils.py @@ -0,0 +1,390 @@ +"""Unit tests for search utility functions.""" + +import pytest + +from upstash_redis.search_utils import ( + flatten_schema, + build_create_index_command, + build_query_command, + deserialize_query_response, + deserialize_describe_response, + parse_count_response, +) + + +class TestFlattenSchema: + """Tests for flatten_schema function.""" + + def test_flatten_simple_schema(self): + """Test flattening a simple schema.""" + schema = {"name": "TEXT", "age": {"type": "U64", "fast": True}, "active": "BOOL"} + + flattened = flatten_schema(schema) + + assert len(flattened) == 3 + assert flattened[0].path == "name" + assert flattened[0].field_type == "TEXT" + assert flattened[1].path == "age" + assert flattened[1].field_type == "U64" + assert flattened[1].fast is True + + def test_flatten_nested_schema(self): + """Test flattening a nested schema.""" + schema = { + "title": "TEXT", + "author": {"name": "TEXT", "email": "TEXT"}, + "stats": {"views": {"type": "U64", "fast": True}}, + } + + flattened = flatten_schema(schema) + + assert len(flattened) == 4 + # Check that nested paths are created correctly + paths = [f.path for f in flattened] + assert "title" in paths + assert "author.name" in paths + assert "author.email" in paths + assert "stats.views" in paths + + def test_flatten_with_field_options(self): + """Test flattening schema with field options.""" + schema = { + "title": {"type": "TEXT", "noStem": True, "noTokenize": True}, + "count": {"type": "U64", "fast": True}, + } + + flattened = flatten_schema(schema) + + title_field = next(f for f in flattened if f.path == "title") + assert title_field.no_stem is True + assert title_field.no_tokenize is True + + +class TestBuildCreateIndexCommand: + """Tests for build_create_index_command function.""" + + def test_build_simple_create_command(self): + """Test building a simple create index command.""" + params = { + "name": "test-idx", + "dataType": "string", + "prefix": "test:", + "schema": {"name": "TEXT", "age": {"type": "U64", "fast": True}}, + } + + command = build_create_index_command(params) + + assert command == [ + "SEARCH.CREATE", + "test-idx", + "ON", + "STRING", + "PREFIX", + "1", + "test:", + "SCHEMA", + "name", + "TEXT", + "age", + "U64", + "FAST", + ] + + def test_build_create_with_language(self): + """Test building create command with language option.""" + params = { + "name": "test-idx", + "dataType": "json", + "prefix": "test:", + "schema": {"content": "TEXT"}, + "language": "turkish", + } + + command = build_create_index_command(params) + + assert command == [ + "SEARCH.CREATE", + "test-idx", + "ON", + "JSON", + "PREFIX", + "1", + "test:", + "LANGUAGE", + "turkish", + "SCHEMA", + "content", + "TEXT", + ] + + def test_build_create_with_multiple_prefixes(self): + """Test building create command with multiple prefixes.""" + params = { + "name": "test-idx", + "dataType": "hash", + "prefix": ["user:", "profile:"], + "schema": {"name": "TEXT"}, + } + + command = build_create_index_command(params) + + assert command == [ + "SEARCH.CREATE", + "test-idx", + "ON", + "HASH", + "PREFIX", + "2", + "user:", + "profile:", + "SCHEMA", + "name", + "TEXT", + ] + + def test_build_create_with_skip_initial_scan(self): + """Test building create command with skipInitialScan.""" + params = { + "name": "test-idx", + "dataType": "string", + "prefix": "test:", + "schema": {"name": "TEXT"}, + "skipInitialScan": True, + } + + command = build_create_index_command(params) + + assert command == [ + "SEARCH.CREATE", + "test-idx", + "SKIPINITIALSCAN", + "ON", + "STRING", + "PREFIX", + "1", + "test:", + "SCHEMA", + "name", + "TEXT", + ] + + def test_build_create_with_field_options(self): + """Test building create command with field options.""" + params = { + "name": "test-idx", + "dataType": "string", + "prefix": "test:", + "schema": { + "title": {"type": "TEXT", "noStem": True, "noTokenize": True}, + "score": {"type": "F64", "fast": True}, + }, + } + + command = build_create_index_command(params) + + assert command == [ + "SEARCH.CREATE", + "test-idx", + "ON", + "STRING", + "PREFIX", + "1", + "test:", + "SCHEMA", + "title", + "TEXT", + "NOTOKENIZE", + "NOSTEM", + "score", + "F64", + "FAST", + ] + + +class TestBuildQueryCommand: + """Tests for build_query_command function.""" + + def test_build_simple_query(self): + """Test building a simple query command.""" + command = build_query_command( + "SEARCH.QUERY", "test-idx", {"filter": {"name": {"$eq": "test"}}} + ) + + assert command[0] == "SEARCH.QUERY" + assert command[1] == "test-idx" + assert '{"name":{"$eq":"test"}}' in command[2] or '{"name": {"$eq": "test"}}' in command[2] + + def test_build_query_with_limit(self): + """Test building query with limit.""" + command = build_query_command( + "SEARCH.QUERY", "test-idx", {"filter": {"name": {"$eq": "test"}}, "limit": 10} + ) + + assert command == [ + "SEARCH.QUERY", + "test-idx", + '{"name":{"$eq":"test"}}', + "LIMIT", + "10", + ] + + def test_build_query_with_offset(self): + """Test building query with offset.""" + command = build_query_command( + "SEARCH.QUERY", "test-idx", {"filter": {"name": {"$eq": "test"}}, "limit": 10, "offset": 5} + ) + + assert command == [ + "SEARCH.QUERY", + "test-idx", + '{"name":{"$eq":"test"}}', + "LIMIT", + "10", + "OFFSET", + "5", + ] + + def test_build_query_with_sorting(self): + """Test building query with sorting.""" + command = build_query_command( + "SEARCH.QUERY", + "test-idx", + {"filter": {"name": {"$eq": "test"}}, "orderBy": {"score": "DESC"}}, + ) + + assert command == [ + "SEARCH.QUERY", + "test-idx", + '{"name":{"$eq":"test"}}', + "SORTBY", + "score", + "DESC", + ] + + def test_build_query_with_select_nocontent(self): + """Test building query with noContent.""" + command = build_query_command( + "SEARCH.QUERY", "test-idx", {"filter": {"name": {"$eq": "test"}}, "select": {}} + ) + + assert command == [ + "SEARCH.QUERY", + "test-idx", + '{"name":{"$eq":"test"}}', + "NOCONTENT", + ] + + def test_build_query_with_select_fields(self): + """Test building query with specific return fields.""" + command = build_query_command( + "SEARCH.QUERY", + "test-idx", + {"filter": {"name": {"$eq": "test"}}, "select": {"name": True, "age": True}}, + ) + + assert command == [ + "SEARCH.QUERY", + "test-idx", + '{"name":{"$eq":"test"}}', + "RETURN", + "2", + "name", + "age", + ] + + def test_build_query_numeric_filters(self): + """Test building query with numeric filters.""" + # Test JSON format + command = build_query_command("SEARCH.QUERY", "test-idx", {"filter": {"price": {"$gt": 100}}}) + assert command == [ + "SEARCH.QUERY", + "test-idx", + '{"price":{"$gt":100}}', + ] + + def test_build_query_text_filters(self): + """Test building query with text filters.""" + # Test JSON format + command = build_query_command( + "SEARCH.QUERY", "test-idx", {"filter": {"name": {"$fuzzy": "test"}}} + ) + assert command == [ + "SEARCH.QUERY", + "test-idx", + '{"name":{"$fuzzy":"test"}}', + ] + + +class TestDeserializeQueryResponse: + """Tests for deserialize_query_response function.""" + + def test_deserialize_simple_response(self): + """Test deserializing a simple query response.""" + raw = [["key1", "0.5", [["name", "test"], ["age", 25]]]] + + result = deserialize_query_response(raw) + + assert len(result) == 1 + assert result[0]["key"] == "key1" + assert result[0]["score"] == "0.5" + assert "data" in result[0] + + def test_deserialize_nocontent_response(self): + """Test deserializing a response with no content.""" + raw = [["key1", "0.5"]] + + result = deserialize_query_response(raw) + + assert len(result) == 1 + assert result[0]["key"] == "key1" + assert result[0]["score"] == "0.5" + assert "data" not in result[0] + + def test_deserialize_nested_fields(self): + """Test deserializing response with nested fields.""" + raw = [["key1", "0.5", [["author.name", "John"], ["author.email", "john@example.com"]]]] + + result = deserialize_query_response(raw) + + assert len(result) == 1 + assert "data" in result[0] + # The deserialization should create nested structure + # (implementation may vary) + + +class TestDeserializeDescribeResponse: + """Tests for deserialize_describe_response function.""" + + def test_deserialize_describe(self): + """Test deserializing a describe response.""" + raw = [ + "name", + "test-idx", + "type", + "STRING", + "prefixes", + ["test:"], + "schema", + [["name", "TEXT"], ["age", "U64", "FAST"]], + ] + + result = deserialize_describe_response(raw) + + assert result["name"] == "test-idx" + assert result["dataType"] == "string" + assert result["prefixes"] == ["test:"] + assert "schema" in result + assert "name" in result["schema"] + assert result["schema"]["name"]["type"] == "TEXT" + + +class TestParseCountResponse: + """Tests for parse_count_response function.""" + + def test_parse_int_response(self): + """Test parsing integer count response.""" + assert parse_count_response(42) == 42 + + def test_parse_string_response(self): + """Test parsing string count response.""" + assert parse_count_response("42") == 42 diff --git a/upstash_redis/asyncio/client.py b/upstash_redis/asyncio/client.py index f805a39..2fcfaf1 100644 --- a/upstash_redis/asyncio/client.py +++ b/upstash_redis/asyncio/client.py @@ -10,6 +10,7 @@ from upstash_redis.format import cast_response from upstash_redis.http import make_headers, AsyncHttpClient from upstash_redis.typing import RESTResultT +from upstash_redis.search_namespace import SearchNamespace class Redis(AsyncCommands): @@ -60,6 +61,7 @@ def __init__( self._headers = make_headers(token, rest_encoding, allow_telemetry) self._json = AsyncJsonCommands(self) + self._search = SearchNamespace(self) self._http = AsyncHttpClient( encoding=rest_encoding, retries=rest_retries, @@ -71,6 +73,10 @@ def __init__( def json(self) -> AsyncJsonCommands: return self._json + @property + def search(self) -> SearchNamespace: + return self._search + @classmethod def from_env( cls, diff --git a/upstash_redis/client.py b/upstash_redis/client.py index c2a9bc8..76b473f 100644 --- a/upstash_redis/client.py +++ b/upstash_redis/client.py @@ -10,6 +10,7 @@ from upstash_redis.format import cast_response from upstash_redis.http import make_headers, SyncHttpClient from upstash_redis.typing import RESTResultT +from upstash_redis.search_namespace import SearchNamespace class Redis(Commands): @@ -62,6 +63,7 @@ def __init__( self._headers = make_headers(token, rest_encoding, allow_telemetry) self._json = JsonCommands(self) + self._search = SearchNamespace(self) self._http = SyncHttpClient( encoding=rest_encoding, retries=rest_retries, @@ -73,6 +75,10 @@ def __init__( def json(self) -> JsonCommands: return self._json + @property + def search(self) -> SearchNamespace: + return self._search + @classmethod def from_env( cls, diff --git a/upstash_redis/search_index.py b/upstash_redis/search_index.py new file mode 100644 index 0000000..8f04315 --- /dev/null +++ b/upstash_redis/search_index.py @@ -0,0 +1,281 @@ +"""Search index functionality for Upstash Redis.""" + +from typing import Any, Dict, List, Optional, Union + +from upstash_redis.search_types import ( + CreateIndexParams, + IndexDescription, + QueryOptions, + QueryResult, + NestedIndexSchema, + FlatIndexSchema, + RootQueryFilter, +) +from upstash_redis.search_utils import ( + build_create_index_command, + build_query_command, + deserialize_describe_response, + deserialize_query_response, + parse_count_response, +) + + +class SearchIndex: + """ + Represents a search index in Upstash Redis. + + This class provides methods to interact with a search index including + querying, counting, describing, and dropping the index. + """ + + def __init__( + self, + name: str, + client: Any, + schema: Optional[Union[NestedIndexSchema, FlatIndexSchema]] = None, + ): + """ + Initialize a SearchIndex instance. + + Args: + name: Name of the index + client: Redis client instance (sync or async) + schema: Optional schema definition for the index + """ + self.name = name + self.schema = schema + self._client = client + + def wait_indexing(self) -> Any: + """ + Wait for the index to finish indexing all existing keys. + + Returns: + The result of the command execution + + Example: + ```python + index.wait_indexing() + ``` + """ + command = ["SEARCH.WAITINDEXING", self.name] + return self._client.execute(command) + + def describe(self) -> IndexDescription: + """ + Get detailed information about the index structure. + + Returns: + IndexDescription with schema, prefixes, and other metadata + + Example: + ```python + description = index.describe() + print(description["schema"]) + ``` + """ + command = ["SEARCH.DESCRIBE", self.name] + raw_result = self._client.execute(command) + + # If the client returns a coroutine (async), we can't deserialize here + # The async version will need to override this method + if hasattr(raw_result, "__await__"): + # For async client, return the coroutine wrapped in deserializer + async def _async_describe(): + result = await raw_result + return deserialize_describe_response(result) + + return _async_describe() + + return deserialize_describe_response(raw_result) + + def query( + self, + *, + filter: RootQueryFilter, + limit: Optional[int] = None, + offset: Optional[int] = None, + orderBy: Optional[Dict[str, str]] = None, + select: Optional[Dict[str, bool]] = None, + highlight: Optional[Dict[str, Any]] = None, + ) -> List[QueryResult]: + """ + Query the index with filters and options. + + Args: + filter: Filter specification mapping field names to conditions + limit: Maximum number of results to return + offset: Number of results to skip (for pagination) + orderBy: Field to sort by with direction ("ASC" or "DESC") + select: Fields to include/exclude in results + highlight: Highlighting options for matched terms + + Returns: + List of query results with keys, scores, and data + + Example: + ```python + # Simple text search + results = index.query(filter={"name": {"$eq": "Laptop"}}) + + # With pagination and sorting + results = index.query( + filter={"category": {"$eq": "electronics"}}, + limit=10, + offset=0, + orderBy={"price": "ASC"} + ) + ``` + """ + options: QueryOptions = {"filter": filter} + if limit is not None: + options["limit"] = limit + if offset is not None: + options["offset"] = offset + if orderBy is not None: + options["orderBy"] = orderBy + if select is not None: + options["select"] = select + if highlight is not None: + options["highlight"] = highlight + + command = build_query_command("SEARCH.QUERY", self.name, options) + raw_result = self._client.execute(command) + + # Handle async + if hasattr(raw_result, "__await__"): + + async def _async_query(): + result = await raw_result + return deserialize_query_response(result) + + return _async_query() + + return deserialize_query_response(raw_result) + + def count(self, filter: RootQueryFilter) -> Dict[str, int]: + """ + Count documents matching a filter. + + Args: + filter: Filter specification mapping field names to conditions + + Returns: + Dictionary with "count" key containing the number of matches + + Example: + ```python + result = index.count({"active": {"$eq": True}}) + print(result["count"]) # e.g., 42 + ``` + """ + command = build_query_command("SEARCH.COUNT", self.name, {"filter": filter}) + raw_result = self._client.execute(command) + + # Handle async + if hasattr(raw_result, "__await__"): + + async def _async_count(): + result = await raw_result + return {"count": parse_count_response(result)} + + return _async_count() + + return {"count": parse_count_response(raw_result)} + + def drop(self) -> Union[int, Any]: + """ + Drop (delete) the index. + + Returns: + 1 if the index was dropped, 0 if it didn't exist + + Example: + ```python + result = index.drop() + print(result) # 1 + ``` + """ + command = ["SEARCH.DROP", self.name] + return self._client.execute(command) + + +def create_index(client: Any, params: CreateIndexParams) -> SearchIndex: + """ + Create a new search index. + + Args: + client: Redis client instance (sync or async) + params: Index creation parameters including name, schema, data type, etc. + + Returns: + SearchIndex instance for the newly created index + + Example: + ```python + schema = { + "name": "TEXT", + "age": {"type": "U64", "fast": True}, + "active": "BOOL" + } + + index = create_index(redis, { + "name": "users-idx", + "prefix": "user:", + "dataType": "json", + "schema": schema + }) + ``` + """ + name = params["name"] + schema = params.get("schema") + + # Build and execute create command + command = build_create_index_command(params) + result = client.execute(command) + + # Handle async + if hasattr(result, "__await__"): + + async def _async_create(): + await result + return SearchIndex(name=name, client=client, schema=schema) + + return _async_create() + + return SearchIndex(name=name, client=client, schema=schema) + + +def init_index( + client: Any, + name: str, + schema: Optional[Union[NestedIndexSchema, FlatIndexSchema]] = None, +) -> SearchIndex: + """ + Initialize a SearchIndex instance for an existing index. + + This function does not create the index, it only creates a Python object + to interact with an existing index. + + Args: + client: Redis client instance (sync or async) + name: Name of the existing index + schema: Optional schema definition for better type hints + + Returns: + SearchIndex instance + + Example: + ```python + # Without schema + index = init_index(redis, "users-idx") + + # With schema for better typing + schema = { + "name": "TEXT", + "age": {"type": "U64", "fast": True} + } + index = init_index(redis, "users-idx", schema) + ``` + """ + return SearchIndex(name=name, client=client, schema=schema) diff --git a/upstash_redis/search_namespace.py b/upstash_redis/search_namespace.py new file mode 100644 index 0000000..552de78 --- /dev/null +++ b/upstash_redis/search_namespace.py @@ -0,0 +1,113 @@ +"""Search namespace for Redis client.""" + +from typing import Any, Optional, Union + +from upstash_redis.search_index import SearchIndex, create_index as _create_index, init_index as _init_index +from upstash_redis.search_types import CreateIndexParams, NestedIndexSchema, FlatIndexSchema + + +class SearchNamespace: + """ + Namespace for search index operations. + + Access via redis.search.createIndex() or redis.search.index() + """ + + def __init__(self, client: Any): + """ + Initialize the search namespace. + + Args: + client: Redis client instance (sync or async) + """ + self._client = client + + def createIndex( + self, + *, + name: str, + schema: Union[NestedIndexSchema, FlatIndexSchema], + dataType: str, + prefix: Union[str, list[str]], + language: Optional[str] = None, + skipInitialScan: bool = False, + existsOk: bool = False, + ) -> SearchIndex: + """ + Create a new search index. + + Args: + name: Name of the index + schema: Schema definition mapping field names to types + dataType: Type of data being indexed ("string", "json", or "hash") + prefix: Key prefix(es) for automatic indexing + language: Optional language for text analysis + skipInitialScan: Skip indexing existing keys + existsOk: Don't error if index already exists + + Returns: + SearchIndex instance for the newly created index + + Example: + ```python + schema = { + "name": "TEXT", + "age": {"type": "U64", "fast": True}, + "active": "BOOL" + } + + index = redis.search.createIndex( + name="users-idx", + prefix="user:", + dataType="json", + schema=schema + ) + ``` + """ + params: CreateIndexParams = { + "name": name, + "schema": schema, + "dataType": dataType, + "prefix": prefix, + } + if language is not None: + params["language"] = language + if skipInitialScan: + params["skipInitialScan"] = skipInitialScan + if existsOk: + params["existsOk"] = existsOk + + return _create_index(self._client, params) + + def index( + self, + name: str, + schema: Optional[Union[NestedIndexSchema, FlatIndexSchema]] = None + ) -> SearchIndex: + """ + Initialize a SearchIndex instance for an existing index. + + This method does not create the index, it only creates a Python object + to interact with an existing index. + + Args: + name: Name of the existing index + schema: Optional schema definition for better type hints + + Returns: + SearchIndex instance + + Example: + ```python + # Without schema + index = redis.search.index("users-idx") + + # With schema for better typing + schema = { + "name": "TEXT", + "age": {"type": "U64", "fast": True} + } + index = redis.search.index("users-idx", schema) + ``` + """ + return _init_index(self._client, name, schema) diff --git a/upstash_redis/search_types.py b/upstash_redis/search_types.py new file mode 100644 index 0000000..c685745 --- /dev/null +++ b/upstash_redis/search_types.py @@ -0,0 +1,158 @@ +"""Type definitions for search index functionality.""" + +from typing import Any, Dict, List, Literal, Optional, TypedDict, Union + +# Field types +FieldType = Literal["TEXT", "U64", "I64", "F64", "BOOL", "DATE"] + +# Language options for text indexing +Language = Literal[ + "arabic", + "basque", + "catalan", + "danish", + "dutch", + "english", + "finnish", + "french", + "german", + "greek", + "hungarian", + "indonesian", + "irish", + "italian", + "lithuanian", + "nepali", + "norwegian", + "portuguese", + "romanian", + "russian", + "spanish", + "swedish", + "tamil", + "turkish", +] + +# Data types that can be indexed +DataType = Literal["string", "json", "hash"] + + +# Schema definitions +class DetailedField(TypedDict, total=False): + """Detailed field configuration with options.""" + + type: FieldType + fast: bool + noTokenize: bool + noStem: bool + from_: str # Using from_ to avoid Python keyword + + +# Schema can be a nested structure of fields +NestedIndexSchema = Dict[str, Union[FieldType, DetailedField, "NestedIndexSchema"]] +FlatIndexSchema = Dict[str, Union[FieldType, DetailedField]] + + +# Query filter operators +class FuzzyOperator(TypedDict, total=False): + """Fuzzy search with typo tolerance.""" + + value: str + distance: int + + +class PhraseOperator(TypedDict, total=False): + """Exact phrase matching.""" + + value: str + prefix: bool + + +# Field filters - using MongoDB-style $ prefix for operators +class TextFieldFilter(TypedDict, total=False): + """Text field query filters.""" + + # Note: TypedDict doesn't support keys starting with $ + # These are represented as string keys at runtime + # Use as: {"$eq": "value"} in actual code + + +class NumericFieldFilter(TypedDict, total=False): + """Numeric field query filters.""" + + # Note: TypedDict doesn't support keys starting with $ + # These are represented as string keys at runtime + # Use as: {"$eq": 5, "$gt": 10} in actual code + + +class BooleanFieldFilter(TypedDict, total=False): + """Boolean field query filters.""" + + # Note: TypedDict doesn't support keys starting with $ + # These are represented as string keys at runtime + # Use as: {"$eq": True} in actual code + + +# Root query filter - maps field names to their filters +RootQueryFilter = Dict[str, Union[TextFieldFilter, NumericFieldFilter, BooleanFieldFilter]] + + +# Query options +class HighlightOptions(TypedDict, total=False): + """Highlighting options for query results.""" + + fields: List[str] + tags: Dict[str, str] # "open" and "close" tags + + +class QueryOptions(TypedDict, total=False): + """Options for querying an index.""" + + filter: RootQueryFilter + limit: int + offset: int + orderBy: Dict[str, Literal["ASC", "DESC"]] + select: Dict[str, bool] # Field name to include/exclude + highlight: HighlightOptions + + +# Query result +class QueryResult(TypedDict, total=False): + """Result from a query operation.""" + + key: str + score: str + data: Dict[str, Any] + + +# Index description +class DescribeFieldInfo(TypedDict, total=False): + """Field information from describe command.""" + + type: FieldType + fast: bool + noTokenize: bool + noStem: bool + + +class IndexDescription(TypedDict, total=False): + """Description of an index.""" + + name: str + dataType: DataType + prefixes: List[str] + language: Language + schema: Dict[str, DescribeFieldInfo] + + +# Create index parameters +class CreateIndexParams(TypedDict, total=False): + """Parameters for creating a new index.""" + + name: str + prefix: Union[str, List[str]] + dataType: DataType + schema: Union[NestedIndexSchema, FlatIndexSchema] + language: Language + skipInitialScan: bool + existsOk: bool diff --git a/upstash_redis/search_utils.py b/upstash_redis/search_utils.py new file mode 100644 index 0000000..95837e3 --- /dev/null +++ b/upstash_redis/search_utils.py @@ -0,0 +1,440 @@ +"""Utility functions for search index operations.""" + +from typing import Any, Dict, List, Union + +from upstash_redis.search_types import ( + CreateIndexParams, + DescribeFieldInfo, + DetailedField, + FieldType, + FlatIndexSchema, + IndexDescription, + NestedIndexSchema, + QueryOptions, + QueryResult, + TextFieldFilter, + NumericFieldFilter, + BooleanFieldFilter, +) + + +class FlattenedField: + """Represents a flattened field from a nested schema.""" + + def __init__( + self, + path: str, + field_type: FieldType, + fast: bool = False, + no_tokenize: bool = False, + no_stem: bool = False, + from_: str = None, + ): + self.path = path + self.field_type = field_type + self.fast = fast + self.no_tokenize = no_tokenize + self.no_stem = no_stem + self.from_ = from_ + + +def _is_field_type(value: Any) -> bool: + """Check if value is a valid field type.""" + return isinstance(value, str) and value in ["TEXT", "U64", "I64", "F64", "BOOL", "DATE"] + + +def _is_detailed_field(value: Any) -> bool: + """Check if value is a detailed field configuration.""" + return isinstance(value, dict) and "type" in value and _is_field_type(value.get("type")) + + +def _is_nested_schema(value: Any) -> bool: + """Check if value is a nested schema.""" + return isinstance(value, dict) and not _is_detailed_field(value) and not _is_field_type(value) + + +def flatten_schema( + schema: Union[NestedIndexSchema, FlatIndexSchema], path_prefix: List[str] = None +) -> List[FlattenedField]: + """ + Flatten a nested schema into a list of field definitions. + + Args: + schema: The schema to flatten + path_prefix: Current path prefix for nested fields + + Returns: + List of flattened field definitions + """ + if path_prefix is None: + path_prefix = [] + + fields: List[FlattenedField] = [] + + for key, value in schema.items(): + current_path = path_prefix + [key] + path_string = ".".join(current_path) + + if _is_field_type(value): + fields.append(FlattenedField(path=path_string, field_type=value)) + elif _is_detailed_field(value): + detailed = value + fields.append( + FlattenedField( + path=path_string, + field_type=detailed["type"], + fast=detailed.get("fast", False), + no_tokenize=detailed.get("noTokenize", False), + no_stem=detailed.get("noStem", False), + from_=detailed.get("from_"), + ) + ) + elif _is_nested_schema(value): + nested_fields = flatten_schema(value, current_path) + fields.extend(nested_fields) + + return fields + + +def build_create_index_command(params: CreateIndexParams) -> List[str]: + """ + Build the SEARCH.CREATE command from parameters. + + Args: + params: Create index parameters + + Returns: + Command list ready to execute + """ + command: List[str] = ["SEARCH.CREATE", params["name"]] + + # Add options BEFORE "ON" keyword + if params.get("skipInitialScan", False): + command.append("SKIPINITIALSCAN") + + if params.get("existsOk", False): + command.append("EXISTSOK") + + # Add "ON" keyword and data type + command.append("ON") + data_type = params["dataType"].upper() + command.append(data_type) + + # Add prefixes + prefixes = params["prefix"] + if isinstance(prefixes, str): + prefixes = [prefixes] + + command.append("PREFIX") + command.append(str(len(prefixes))) + command.extend(prefixes) + + # Add language if specified + if "language" in params: + command.extend(["LANGUAGE", params["language"]]) + + # Add schema + schema = params["schema"] + flattened = flatten_schema(schema) + + command.append("SCHEMA") + for field in flattened: + # Add field path or from source + if field.from_: + command.extend(["AS", field.path, "FROM", field.from_]) + else: + command.append(field.path) + + # Add field type + command.append(field.field_type) + + # Add field options + if field.no_tokenize: + command.append("NOTOKENIZE") + if field.no_stem: + command.append("NOSTEM") + if field.fast: + command.append("FAST") + + return command + + +def _build_filter_query(filter_dict: Dict[str, Any]) -> List[str]: + """ + Build filter query from filter dictionary. + + Args: + filter_dict: Filter specification + + Returns: + List of filter query parts + """ + filters: List[str] = [] + + for field_name, field_filter in filter_dict.items(): + if not isinstance(field_filter, dict): + continue + + # Text filters + if "eq" in field_filter: + value = field_filter["eq"] + if isinstance(value, bool): + filters.append(f"@{field_name}:[{str(value).lower()}]") + else: + filters.append(f"@{field_name}:({value})") + + elif "fuzzy" in field_filter: + fuzzy = field_filter["fuzzy"] + if isinstance(fuzzy, str): + filters.append(f"@{field_name}:%{fuzzy}%") + elif isinstance(fuzzy, dict): + value = fuzzy.get("value", "") + distance = fuzzy.get("distance", 1) + filters.append(f"@{field_name}:%" + "%" * distance + value + "%" * distance) + + elif "phrase" in field_filter: + phrase = field_filter["phrase"] + if isinstance(phrase, str): + filters.append(f'@{field_name}:"{phrase}"') + elif isinstance(phrase, dict): + value = phrase.get("value", "") + prefix = phrase.get("prefix", False) + if prefix: + filters.append(f'@{field_name}:"{value}"*') + else: + filters.append(f'@{field_name}:"{value}"') + + elif "regex" in field_filter: + regex = field_filter["regex"] + filters.append(f"@{field_name}:/{regex}/") + + # Numeric filters + elif "gt" in field_filter or "gte" in field_filter or "lt" in field_filter or "lte" in field_filter: + min_val = "-inf" + max_val = "+inf" + + if "gt" in field_filter: + min_val = f"({field_filter['gt']}" + elif "gte" in field_filter: + min_val = str(field_filter["gte"]) + + if "lt" in field_filter: + max_val = f"({field_filter['lt']}" + elif "lte" in field_filter: + max_val = str(field_filter["lte"]) + + filters.append(f"@{field_name}:[{min_val} {max_val}]") + + return filters + + +def build_query_command( + command_name: str, index_name: str, options: QueryOptions = None +) -> List[str]: + """ + Build SEARCH.QUERY or SEARCH.COUNT command. + + Args: + command_name: "SEARCH.QUERY" or "SEARCH.COUNT" + index_name: Name of the index + options: Query options + + Returns: + Command list ready to execute + """ + import json + + command: List[str] = [command_name, index_name] + + # Serialize filter to JSON (compact format without spaces) + filter_json = "{}" + if options and "filter" in options: + filter_json = json.dumps(options["filter"], separators=(',', ':')) + + command.append(filter_json) + + if not options: + return command + + # Add limit + if "limit" in options: + command.extend(["LIMIT", str(options["limit"])]) + + # Add offset + if "offset" in options: + command.extend(["OFFSET", str(options["offset"])]) + + # Add orderBy + if "orderBy" in options: + for field, direction in options["orderBy"].items(): + command.extend(["SORTBY", field, direction]) + + # Add highlight + if "highlight" in options: + highlight = options["highlight"] + command.append("HIGHLIGHT") + if "fields" in highlight: + fields = highlight["fields"] + command.extend(["FIELDS", str(len(fields))]) + command.extend(fields) + if "preTag" in highlight and "postTag" in highlight: + command.extend(["TAGS", highlight["preTag"], highlight["postTag"]]) + + # Add select + if "select" in options: + select = options["select"] + if not select: # Empty dict means NOCONTENT + command.append("NOCONTENT") + else: + command.append("RETURN") + fields = [f for f, include in select.items() if include] + command.append(str(len(fields))) + command.extend(fields) + + return command + + # Add highlight + if "highlight" in options: + highlight = options["highlight"] + command.append("HIGHLIGHT") + + if "fields" in highlight and highlight["fields"]: + command.append("FIELDS") + command.append(str(len(highlight["fields"]))) + command.extend(highlight["fields"]) + + if "tags" in highlight: + tags = highlight["tags"] + if "open" in tags and "close" in tags: + command.extend(["TAGS", tags["open"], tags["close"]]) + + return command + + +def deserialize_query_response(raw_response: List[Any]) -> List[QueryResult]: + """ + Deserialize raw query response into structured results. + + Args: + raw_response: Raw response from SEARCH.QUERY + + Returns: + List of query results + """ + results: List[QueryResult] = [] + + for item_raw in raw_response: + if not isinstance(item_raw, (list, tuple)) or len(item_raw) < 2: + continue + + key = item_raw[0] + score = item_raw[1] + + result: QueryResult = {"key": key, "score": score} + + if len(item_raw) > 2: + raw_fields = item_raw[2] + + if isinstance(raw_fields, (list, tuple)) and len(raw_fields) > 0: + data: Dict[str, Any] = {} + + # Process field pairs + for field_raw in raw_fields: + if not isinstance(field_raw, (list, tuple)) or len(field_raw) < 2: + continue + + field_key = field_raw[0] + field_value = field_raw[1] + + # Handle nested paths + path_parts = field_key.split(".") + if len(path_parts) == 1: + data[field_key] = field_value + else: + # Build nested structure + current_obj = data + for i, part in enumerate(path_parts[:-1]): + if part not in current_obj: + current_obj[part] = {} + current_obj = current_obj[part] + current_obj[path_parts[-1]] = field_value + + # If $ key exists (full document), use its contents + if "$" in data: + data = data["$"] + + result["data"] = data + + results.append(result) + + return results + + +def deserialize_describe_response(raw_response: List[Any]) -> IndexDescription: + """ + Deserialize raw describe response into index description. + + Args: + raw_response: Raw response from SEARCH.DESCRIBE + + Returns: + Index description + """ + description: IndexDescription = {} + + i = 0 + while i < len(raw_response): + descriptor = raw_response[i] + value = raw_response[i + 1] if i + 1 < len(raw_response) else None + + if descriptor == "name": + description["name"] = value + elif descriptor == "type": + if value: + description["dataType"] = value.lower() + elif descriptor == "prefixes": + description["prefixes"] = value if isinstance(value, list) else [] + elif descriptor == "language": + description["language"] = value + elif descriptor == "schema": + schema: Dict[str, DescribeFieldInfo] = {} + if isinstance(value, list): + for field_desc in value: + if not isinstance(field_desc, list) or len(field_desc) < 2: + continue + + field_name = field_desc[0] + field_info: DescribeFieldInfo = {"type": field_desc[1]} + + # Parse field options + for j in range(2, len(field_desc)): + option = field_desc[j] + if option == "NOSTEM": + field_info["noStem"] = True + elif option == "NOTOKENIZE": + field_info["noTokenize"] = True + elif option == "FAST": + field_info["fast"] = True + + schema[field_name] = field_info + + description["schema"] = schema + + i += 2 + + return description + + +def parse_count_response(raw_response: Any) -> int: + """ + Parse count response. + + Args: + raw_response: Raw response from SEARCH.COUNT + + Returns: + Count as integer + """ + if isinstance(raw_response, int): + return raw_response + return int(raw_response) From e9eb4e81456843e90359c5838dced3c69edee2ec Mon Sep 17 00:00:00 2001 From: CahidArda Date: Wed, 21 Jan 2026 12:46:53 +0300 Subject: [PATCH 02/12] fix: replace createIndex with create_index, rm unused types, add types to tests --- tests/commands/search/test_search_advanced.py | 10 ++-- tests/commands/search/test_search_index.py | 46 +++++++++--------- upstash_redis/search_index.py | 4 +- upstash_redis/search_namespace.py | 8 ++-- upstash_redis/search_types.py | 47 +------------------ upstash_redis/search_utils.py | 4 -- 6 files changed, 36 insertions(+), 83 deletions(-) diff --git a/tests/commands/search/test_search_advanced.py b/tests/commands/search/test_search_advanced.py index 6ad8cea..c444826 100644 --- a/tests/commands/search/test_search_advanced.py +++ b/tests/commands/search/test_search_advanced.py @@ -26,7 +26,7 @@ def random_id() -> str: @pytest.fixture -def redis_client(): +def redis_client() -> Redis: """Create a Redis client for testing.""" return Redis.from_env() @@ -52,7 +52,7 @@ def json_index(self) -> Generator[IndexFixture, None, None]: "active": "BOOL", } - index = redis.search.createIndex( + index = redis.search.create_index( name=name, schema=schema, dataType="json", prefix=prefix ) @@ -170,7 +170,7 @@ def hash_index(self) -> Generator[IndexFixture, None, None]: schema = {"name": "TEXT", "score": {"type": "U64", "fast": True}} - index = redis.search.createIndex( + index = redis.search.create_index( name=name, schema=schema, dataType="hash", prefix=prefix ) @@ -231,7 +231,7 @@ def nested_string_index(self) -> Generator[IndexFixture, None, None]: "stats": {"views": {"type": "U64", "fast": True}, "likes": {"type": "U64", "fast": True}}, } - index = redis.search.createIndex( + index = redis.search.create_index( name=name, schema=schema, dataType="string", prefix=prefix ) @@ -314,7 +314,7 @@ def nested_json_index(self) -> Generator[IndexFixture, None, None]: "stats": {"views": {"type": "U64", "fast": True}, "likes": {"type": "U64", "fast": True}}, } - index = redis.search.createIndex( + index = redis.search.create_index( name=name, schema=schema, dataType="json", prefix=prefix ) diff --git a/tests/commands/search/test_search_index.py b/tests/commands/search/test_search_index.py index 29ca883..95a2f95 100644 --- a/tests/commands/search/test_search_index.py +++ b/tests/commands/search/test_search_index.py @@ -35,7 +35,7 @@ def random_id() -> str: @pytest.fixture -def redis_client(): +def redis_client() -> Redis: """Create a Redis client for testing.""" return Redis.from_env() @@ -60,21 +60,21 @@ def cleanup_indexes(): class TestCreateIndex: """Tests for creating search indexes.""" - def test_create_string_index_simple_schema(self, redis_client, cleanup_indexes): + def test_create_string_index_simple_schema(self, redis_client: Redis, cleanup_indexes): """Test creating a string index with simple schema.""" name = f"test-string-{random_id()}" cleanup_indexes.append(name) schema = {"name": "TEXT", "age": {"type": "U64", "fast": True}} - index = redis_client.search.createIndex( + index = redis_client.search.create_index( name=name, schema=schema, dataType="string", prefix=f"{name}:" ) assert index is not None assert index.name == name - def test_create_string_index_nested_schema(self, redis_client, cleanup_indexes): + def test_create_string_index_nested_schema(self, redis_client: Redis, cleanup_indexes): """Test creating a string index with nested schema.""" name = f"test-string-nested-{random_id()}" cleanup_indexes.append(name) @@ -84,40 +84,40 @@ def test_create_string_index_nested_schema(self, redis_client, cleanup_indexes): "metadata": {"author": "TEXT", "views": {"type": "U64", "fast": True}}, } - index = redis_client.search.createIndex( + index = redis_client.search.create_index( name=name, schema=schema, dataType="string", prefix=f"{name}:" ) assert index is not None - def test_create_hash_index(self, redis_client, cleanup_indexes): + def test_create_hash_index(self, redis_client: Redis, cleanup_indexes): """Test creating a hash index.""" name = f"test-hash-{random_id()}" cleanup_indexes.append(name) schema = {"title": "TEXT", "count": {"type": "U64", "fast": True}} - index = redis_client.search.createIndex( + index = redis_client.search.create_index( name=name, schema=schema, dataType="hash", prefix=f"{name}:" ) assert index is not None - def test_create_json_index_simple_schema(self, redis_client, cleanup_indexes): + def test_create_json_index_simple_schema(self, redis_client: Redis, cleanup_indexes): """Test creating a JSON index with simple schema.""" name = f"test-json-{random_id()}" cleanup_indexes.append(name) schema = {"name": "TEXT", "age": {"type": "U64", "fast": True}} - index = redis_client.search.createIndex( + index = redis_client.search.create_index( name=name, schema=schema, dataType="json", prefix=f"{name}:" ) assert index is not None assert index.name == name - def test_create_json_index_nested_schema(self, redis_client, cleanup_indexes): + def test_create_json_index_nested_schema(self, redis_client: Redis, cleanup_indexes): """Test creating a JSON index with nested schema.""" name = f"test-json-nested-{random_id()}" cleanup_indexes.append(name) @@ -127,20 +127,20 @@ def test_create_json_index_nested_schema(self, redis_client, cleanup_indexes): "metadata": {"author": "TEXT", "views": {"type": "U64", "fast": True}}, } - index = redis_client.search.createIndex( + index = redis_client.search.create_index( name=name, schema=schema, dataType="json", prefix=f"{name}:" ) assert index is not None - def test_create_index_with_language(self, redis_client, cleanup_indexes): + def test_create_index_with_language(self, redis_client: Redis, cleanup_indexes): """Test creating an index with language option.""" name = f"test-lang-{random_id()}" cleanup_indexes.append(name) schema = {"content": "TEXT"} - index = redis_client.search.createIndex( + index = redis_client.search.create_index( name=name, schema=schema, dataType="string", @@ -150,14 +150,14 @@ def test_create_index_with_language(self, redis_client, cleanup_indexes): assert index is not None - def test_create_index_with_multiple_prefixes(self, redis_client, cleanup_indexes): + def test_create_index_with_multiple_prefixes(self, redis_client: Redis, cleanup_indexes): """Test creating an index with multiple prefixes.""" name = f"test-multi-prefix-{random_id()}" cleanup_indexes.append(name) schema = {"name": "TEXT"} - index = redis_client.search.createIndex( + index = redis_client.search.create_index( name=name, schema=schema, dataType="hash", @@ -188,7 +188,7 @@ def string_index(self) -> Generator[StringIndexFixture, None, None]: "active": "BOOL", } - index = redis.search.createIndex( + index = redis.search.create_index( name=name, schema=schema, dataType="string", prefix=prefix ) @@ -394,7 +394,7 @@ def count_index(self) -> Generator[CountIndexFixture, None, None]: schema = {"type": "TEXT", "value": {"type": "U64", "fast": True}} - index = redis.search.createIndex( + index = redis.search.create_index( name=name, schema=schema, dataType="string", prefix=prefix ) @@ -436,7 +436,7 @@ def test_count_with_numeric_filter(self, count_index: CountIndexFixture): class TestDescribe: """Tests for describe functionality.""" - def test_describe_index(self, redis_client, cleanup_indexes): + def test_describe_index(self, redis_client: Redis, cleanup_indexes): """Test describing an index structure.""" name = f"test-describe-{random_id()}" cleanup_indexes.append(name) @@ -447,7 +447,7 @@ def test_describe_index(self, redis_client, cleanup_indexes): "active": "BOOL", } - index = redis_client.search.createIndex( + index = redis_client.search.create_index( name=name, schema=schema, dataType="string", prefix=f"{name}:" ) @@ -460,13 +460,13 @@ def test_describe_index(self, redis_client, cleanup_indexes): class TestDrop: """Tests for drop functionality.""" - def test_drop_existing_index(self, redis_client): + def test_drop_existing_index(self, redis_client: Redis): """Test dropping an existing index.""" name = f"test-drop-{random_id()}" schema = {"name": "TEXT"} - index = redis_client.search.createIndex( + index = redis_client.search.create_index( name=name, schema=schema, dataType="string", prefix=f"{name}:" ) @@ -478,7 +478,7 @@ def test_drop_existing_index(self, redis_client): class TestIndexMethod: """Tests for redis.search.index method.""" - def test_index_without_schema(self, redis_client): + def test_index_without_schema(self, redis_client: Redis): """Test creating a SearchIndex instance without schema.""" index = redis_client.search.index("test-index") @@ -486,7 +486,7 @@ def test_index_without_schema(self, redis_client): assert index.name == "test-index" assert index.schema is None - def test_index_with_schema(self, redis_client): + def test_index_with_schema(self, redis_client: Redis): """Test creating a SearchIndex instance with schema.""" schema = {"name": "TEXT", "age": {"type": "U64", "fast": True}} diff --git a/upstash_redis/search_index.py b/upstash_redis/search_index.py index 8f04315..6a86449 100644 --- a/upstash_redis/search_index.py +++ b/upstash_redis/search_index.py @@ -200,7 +200,7 @@ def drop(self) -> Union[int, Any]: return self._client.execute(command) -def create_index(client: Any, params: CreateIndexParams) -> SearchIndex: +def _create_index(client: Any, params: CreateIndexParams) -> SearchIndex: """ Create a new search index. @@ -246,7 +246,7 @@ async def _async_create(): return SearchIndex(name=name, client=client, schema=schema) -def init_index( +def _init_index( client: Any, name: str, schema: Optional[Union[NestedIndexSchema, FlatIndexSchema]] = None, diff --git a/upstash_redis/search_namespace.py b/upstash_redis/search_namespace.py index 552de78..f89c61b 100644 --- a/upstash_redis/search_namespace.py +++ b/upstash_redis/search_namespace.py @@ -2,7 +2,7 @@ from typing import Any, Optional, Union -from upstash_redis.search_index import SearchIndex, create_index as _create_index, init_index as _init_index +from upstash_redis.search_index import SearchIndex, _create_index, _init_index from upstash_redis.search_types import CreateIndexParams, NestedIndexSchema, FlatIndexSchema @@ -10,7 +10,7 @@ class SearchNamespace: """ Namespace for search index operations. - Access via redis.search.createIndex() or redis.search.index() + Access via redis.search.create_index() or redis.search.index() """ def __init__(self, client: Any): @@ -22,7 +22,7 @@ def __init__(self, client: Any): """ self._client = client - def createIndex( + def create_index( self, *, name: str, @@ -56,7 +56,7 @@ def createIndex( "active": "BOOL" } - index = redis.search.createIndex( + index = redis.search.create_index( name="users-idx", prefix="user:", dataType="json", diff --git a/upstash_redis/search_types.py b/upstash_redis/search_types.py index c685745..4aaae2a 100644 --- a/upstash_redis/search_types.py +++ b/upstash_redis/search_types.py @@ -52,50 +52,7 @@ class DetailedField(TypedDict, total=False): NestedIndexSchema = Dict[str, Union[FieldType, DetailedField, "NestedIndexSchema"]] FlatIndexSchema = Dict[str, Union[FieldType, DetailedField]] - -# Query filter operators -class FuzzyOperator(TypedDict, total=False): - """Fuzzy search with typo tolerance.""" - - value: str - distance: int - - -class PhraseOperator(TypedDict, total=False): - """Exact phrase matching.""" - - value: str - prefix: bool - - -# Field filters - using MongoDB-style $ prefix for operators -class TextFieldFilter(TypedDict, total=False): - """Text field query filters.""" - - # Note: TypedDict doesn't support keys starting with $ - # These are represented as string keys at runtime - # Use as: {"$eq": "value"} in actual code - - -class NumericFieldFilter(TypedDict, total=False): - """Numeric field query filters.""" - - # Note: TypedDict doesn't support keys starting with $ - # These are represented as string keys at runtime - # Use as: {"$eq": 5, "$gt": 10} in actual code - - -class BooleanFieldFilter(TypedDict, total=False): - """Boolean field query filters.""" - - # Note: TypedDict doesn't support keys starting with $ - # These are represented as string keys at runtime - # Use as: {"$eq": True} in actual code - - -# Root query filter - maps field names to their filters -RootQueryFilter = Dict[str, Union[TextFieldFilter, NumericFieldFilter, BooleanFieldFilter]] - +RootQueryFilter = Dict[str, Any] # Query options class HighlightOptions(TypedDict, total=False): @@ -108,7 +65,7 @@ class HighlightOptions(TypedDict, total=False): class QueryOptions(TypedDict, total=False): """Options for querying an index.""" - filter: RootQueryFilter + filter: Any limit: int offset: int orderBy: Dict[str, Literal["ASC", "DESC"]] diff --git a/upstash_redis/search_utils.py b/upstash_redis/search_utils.py index 95837e3..4c32b28 100644 --- a/upstash_redis/search_utils.py +++ b/upstash_redis/search_utils.py @@ -5,16 +5,12 @@ from upstash_redis.search_types import ( CreateIndexParams, DescribeFieldInfo, - DetailedField, FieldType, FlatIndexSchema, IndexDescription, NestedIndexSchema, QueryOptions, QueryResult, - TextFieldFilter, - NumericFieldFilter, - BooleanFieldFilter, ) From 1186d5256580e240827fc13339fbe1205056425d Mon Sep 17 00:00:00 2001 From: CahidArda Date: Wed, 21 Jan 2026 15:02:17 +0300 Subject: [PATCH 03/12] feat: write search index/commands as Commands interfaces --- .../asyncio/search/test_search_index_async.py | 341 ++++++++++++++++++ tests/commands/search/test_search_advanced.py | 4 +- tests/commands/search/test_search_index.py | 6 +- tests/test_formatters.py | 125 +++++++ upstash_redis/asyncio/client.py | 6 +- upstash_redis/client.py | 6 +- upstash_redis/commands.py | 151 ++++++++ upstash_redis/commands.pyi | 73 ++++ upstash_redis/format.py | 109 +++++- upstash_redis/search_index.py | 281 --------------- upstash_redis/search_namespace.py | 113 ------ 11 files changed, 809 insertions(+), 406 deletions(-) create mode 100644 tests/commands/asyncio/search/test_search_index_async.py delete mode 100644 upstash_redis/search_index.py delete mode 100644 upstash_redis/search_namespace.py diff --git a/tests/commands/asyncio/search/test_search_index_async.py b/tests/commands/asyncio/search/test_search_index_async.py new file mode 100644 index 0000000..39713d7 --- /dev/null +++ b/tests/commands/asyncio/search/test_search_index_async.py @@ -0,0 +1,341 @@ +"""Async tests for search index functionality.""" + +import json +import random +import string +from typing import Any, Dict, List + +import pytest +import pytest_asyncio + +from upstash_redis.asyncio import Redis + + +def random_id() -> str: + """Generate a random ID for test isolation.""" + return "".join(random.choices(string.ascii_lowercase + string.digits, k=8)) + + +@pytest_asyncio.fixture +async def async_redis_client(): + """Create an async Redis client for testing.""" + return Redis.from_env() + + +@pytest_asyncio.fixture +async def cleanup_indexes(): + """Track and cleanup indexes created during tests.""" + indexes = [] + + yield indexes + + # Cleanup after test + redis = Redis.from_env() + for index_name in indexes: + try: + index = redis.search.index(index_name) + await index.drop() + except: + pass + + +@pytest.mark.asyncio +class TestAsyncQuery: + """Tests for async query operations.""" + + @pytest_asyncio.fixture + async def query_index(self, async_redis_client, cleanup_indexes): + """Create a test index with data for querying.""" + name = f"test-async-query-{random_id()}" + prefix = f"{name}:" + keys = [] + cleanup_indexes.append(name) + + redis = async_redis_client + + schema = { + "name": "TEXT", + "description": "TEXT", + "price": {"type": "F64", "fast": True}, + "stock": {"type": "U64", "fast": True}, + "active": "BOOL", + } + + index = await redis.search.create_index( + name=name, schema=schema, dataType="string", prefix=prefix + ) + + # Add test data + test_data = [ + { + "name": "Laptop Pro", + "description": "High performance laptop", + "price": 1299.99, + "stock": 50, + "active": True, + }, + { + "name": "Laptop Basic", + "description": "Budget friendly laptop", + "price": 599.99, + "stock": 100, + "active": True, + }, + { + "name": "Wireless Mouse", + "description": "Ergonomic wireless mouse", + "price": 29.99, + "stock": 200, + "active": True, + }, + { + "name": "USB Cable", + "description": "Fast charging USB cable", + "price": 9.99, + "stock": 500, + "active": False, + }, + ] + + for i, datum in enumerate(test_data): + key = f"{prefix}{i}" + keys.append(key) + await redis.set(key, json.dumps(datum)) + + # Wait for indexing + await index.wait_indexing() + + yield {"index": index, "redis": redis, "keys": keys, "name": name} + + # Cleanup + try: + await index.drop() + except: + pass + if keys: + await redis.delete(*keys) + + async def test_query_text_eq(self, query_index): + """Test async querying with $eq on text field.""" + index = query_index["index"] + result = await index.query(filter={"name": {"$eq": "Laptop"}}) + + assert len(result) > 0 + assert all("key" in r and "score" in r for r in result) + + async def test_query_text_fuzzy(self, query_index): + """Test async querying with fuzzy search.""" + index = query_index["index"] + result = await index.query(filter={"name": {"$fuzzy": "laptopp"}}) + + assert len(result) > 0 + + async def test_query_text_phrase(self, query_index): + """Test async querying with phrase matching.""" + index = query_index["index"] + result = await index.query(filter={"description": {"$phrase": "wireless mouse"}}) + + assert len(result) > 0 + + async def test_query_numeric_gt(self, query_index): + """Test async querying with $gt on numeric field.""" + index = query_index["index"] + result = await index.query(filter={"price": {"$gt": 500}}) + + assert len(result) > 0 + # Verify data structure + for item in result: + assert "key" in item + assert "score" in item + + async def test_query_numeric_gte(self, query_index): + """Test async querying with $gte on numeric field.""" + index = query_index["index"] + result = await index.query(filter={"stock": {"$gte": 100}}) + + assert len(result) > 0 + + async def test_query_numeric_lt(self, query_index): + """Test async querying with $lt on numeric field.""" + index = query_index["index"] + result = await index.query(filter={"price": {"$lt": 50}}) + + assert len(result) > 0 + + async def test_query_boolean_eq(self, query_index): + """Test async querying with $eq on boolean field.""" + index = query_index["index"] + result = await index.query(filter={"active": {"$eq": False}}) + + assert len(result) > 0 + + async def test_query_with_limit(self, query_index): + """Test async querying with limit option.""" + index = query_index["index"] + result = await index.query(filter={"active": {"$eq": True}}, limit=2) + + assert len(result) > 0 + assert len(result) <= 2 + + async def test_query_with_pagination(self, query_index): + """Test async querying with offset for pagination.""" + index = query_index["index"] + first_page = await index.query( + filter={"active": {"$eq": True}}, limit=2, offset=0 + ) + second_page = await index.query( + filter={"active": {"$eq": True}}, limit=2, offset=2 + ) + + assert len(first_page) > 0 + + async def test_query_with_sort_asc(self, query_index): + """Test async querying with sortBy ascending.""" + index = query_index["index"] + result = await index.query( + filter={"active": {"$eq": True}}, + orderBy={"price": "ASC"}, + limit=3, + ) + + assert len(result) > 0 + + async def test_query_with_sort_desc(self, query_index): + """Test async querying with sortBy descending.""" + index = query_index["index"] + result = await index.query( + filter={"active": {"$eq": True}}, + orderBy={"price": "DESC"}, + limit=3, + ) + + assert len(result) > 0 + + +@pytest.mark.asyncio +class TestAsyncCount: + """Tests for async count operations.""" + + @pytest_asyncio.fixture + async def count_index(self, async_redis_client, cleanup_indexes): + """Create a test index with data for counting.""" + name = f"test-async-count-{random_id()}" + prefix = f"{name}:" + keys = [] + cleanup_indexes.append(name) + + redis = async_redis_client + + schema = { + "category": {"type": "TEXT", "noTokenize": True}, + "price": {"type": "F64", "fast": True}, + "active": "BOOL", + } + + index = await redis.search.create_index( + name=name, schema=schema, dataType="string", prefix=prefix + ) + + # Add test data + test_data = [ + {"category": "electronics", "price": 1299.99, "active": True}, + {"category": "electronics", "price": 599.99, "active": True}, + {"category": "accessories", "price": 29.99, "active": True}, + {"category": "accessories", "price": 9.99, "active": False}, + ] + + for i, datum in enumerate(test_data): + key = f"{prefix}{i}" + keys.append(key) + await redis.set(key, json.dumps(datum)) + + # Wait for indexing + await index.wait_indexing() + + yield {"index": index, "redis": redis, "keys": keys, "name": name} + + # Cleanup + try: + await index.drop() + except: + pass + if keys: + await redis.delete(*keys) + + async def test_count_matching_documents(self, count_index): + """Test async counting matching documents.""" + index = count_index["index"] + result = await index.count(filter={"category": {"$eq": "electronics"}}) + + assert "count" in result + assert result["count"] >= 2 + + async def test_count_with_numeric_filter(self, count_index): + """Test async counting with numeric filter.""" + index = count_index["index"] + result = await index.count(filter={"price": {"$gt": 100}}) + + assert "count" in result + assert result["count"] >= 2 + + +@pytest.mark.asyncio +class TestAsyncDescribe: + """Tests for async describe operations.""" + + async def test_describe_index(self, async_redis_client, cleanup_indexes): + """Test async describing an index.""" + name = f"test-async-describe-{random_id()}" + cleanup_indexes.append(name) + + redis = async_redis_client + + schema = { + "name": "TEXT", + "price": {"type": "F64", "fast": True}, + "active": "BOOL", + } + + index = await redis.search.create_index( + name=name, schema=schema, dataType="string", prefix=f"{name}:" + ) + + description = await index.describe() + + assert description is not None + assert "name" in description + assert description["name"] == name + assert "schema" in description + assert "name" in description["schema"] + assert "price" in description["schema"] + assert "active" in description["schema"] + + # Cleanup + try: + await index.drop() + except: + pass + + +@pytest.mark.asyncio +class TestAsyncDrop: + """Tests for async drop operations.""" + + async def test_drop_index(self, async_redis_client): + """Test async dropping an index.""" + name = f"test-async-drop-{random_id()}" + + redis = async_redis_client + + schema = {"name": "TEXT"} + + index = await redis.search.create_index( + name=name, schema=schema, dataType="string", prefix=f"{name}:" + ) + + # Drop the index + result = await index.drop() + + # Verify it was dropped (describe should fail or return None) + # For now just check that drop executed + assert result is not None diff --git a/tests/commands/search/test_search_advanced.py b/tests/commands/search/test_search_advanced.py index c444826..2fc407b 100644 --- a/tests/commands/search/test_search_advanced.py +++ b/tests/commands/search/test_search_advanced.py @@ -6,12 +6,12 @@ import pytest from upstash_redis import Redis -from upstash_redis.search_index import SearchIndex +from upstash_redis.commands import SearchIndexCommands class IndexFixture(TypedDict): """Type definition for index fixtures.""" - index: SearchIndex + index: SearchIndexCommands redis: Redis keys: List[str] name: str diff --git a/tests/commands/search/test_search_index.py b/tests/commands/search/test_search_index.py index 95a2f95..e8150d1 100644 --- a/tests/commands/search/test_search_index.py +++ b/tests/commands/search/test_search_index.py @@ -7,12 +7,12 @@ import pytest from upstash_redis import Redis -from upstash_redis.search_index import SearchIndex +from upstash_redis.commands import SearchIndexCommands class StringIndexFixture(TypedDict): """Type definition for string_index fixture.""" - index: SearchIndex + index: SearchIndexCommands redis: Redis keys: List[str] name: str @@ -20,7 +20,7 @@ class StringIndexFixture(TypedDict): class CountIndexFixture(TypedDict): """Type definition for count_index fixture.""" - index: SearchIndex + index: SearchIndexCommands redis: Redis keys: List[str] name: str diff --git a/tests/test_formatters.py b/tests/test_formatters.py index 318fdd7..32b0bb9 100644 --- a/tests/test_formatters.py +++ b/tests/test_formatters.py @@ -7,6 +7,9 @@ string_to_json, to_json_list, to_optional_bool_list, + format_search_query_response, + format_search_describe_response, + format_search_count_response, ) from upstash_redis.utils import GeoSearchResult @@ -170,3 +173,125 @@ def test_list_to_optional_bool_list() -> None: None, True, ] + + +def test_format_search_query_response() -> None: + # Test basic query response + raw_response = [ + ["doc:1", 0.95, [["name", "Laptop"], ["price", "999"]]], + ["doc:2", 0.85, [["name", "Phone"], ["price", "699"]]], + ] + result = format_search_query_response(raw_response, None) + + assert len(result) == 2 + assert result[0]["key"] == "doc:1" + assert result[0]["score"] == 0.95 + assert result[0]["data"]["name"] == "Laptop" + assert result[0]["data"]["price"] == "999" + assert result[1]["key"] == "doc:2" + assert result[1]["score"] == 0.85 + assert result[1]["data"]["name"] == "Phone" + assert result[1]["data"]["price"] == "699" + + +def test_format_search_query_response_with_nested_paths() -> None: + # Test query response with nested paths + raw_response = [ + ["doc:1", 0.95, [["user.name", "John"], ["user.age", "30"]]], + ] + result = format_search_query_response(raw_response, None) + + assert len(result) == 1 + assert result[0]["key"] == "doc:1" + assert result[0]["score"] == 0.95 + assert result[0]["data"]["user"]["name"] == "John" + assert result[0]["data"]["user"]["age"] == "30" + + +def test_format_search_query_response_with_dollar_key() -> None: + # Test query response with $ key (full document) + raw_response = [ + ["doc:1", 1.0, [["$", {"name": "Laptop", "price": 999}]]], + ] + result = format_search_query_response(raw_response, None) + + assert len(result) == 1 + assert result[0]["key"] == "doc:1" + assert result[0]["score"] == 1.0 + assert result[0]["data"]["name"] == "Laptop" + assert result[0]["data"]["price"] == 999 + + +def test_format_search_query_response_without_fields() -> None: + # Test query response without field data + raw_response = [ + ["doc:1", 0.95], + ["doc:2", 0.85], + ] + result = format_search_query_response(raw_response, None) + + assert len(result) == 2 + assert result[0]["key"] == "doc:1" + assert result[0]["score"] == 0.95 + assert "data" not in result[0] + assert result[1]["key"] == "doc:2" + assert result[1]["score"] == 0.85 + assert "data" not in result[1] + + +def test_format_search_describe_response() -> None: + # Test describe response + raw_response = [ + "name", "myindex", + "type", "JSON", + "prefixes", ["product:", "item:"], + "language", "english", + "schema", [ + ["name", "TEXT", "NOSTEM"], + ["price", "F64", "FAST"], + ["active", "BOOL"], + ], + ] + result = format_search_describe_response(raw_response, None) + + assert result["name"] == "myindex" + assert result["dataType"] == "json" + assert result["prefixes"] == ["product:", "item:"] + assert result["language"] == "english" + assert "name" in result["schema"] + assert result["schema"]["name"]["type"] == "TEXT" + assert result["schema"]["name"]["noStem"] is True + assert "price" in result["schema"] + assert result["schema"]["price"]["type"] == "F64" + assert result["schema"]["price"]["fast"] is True + assert "active" in result["schema"] + assert result["schema"]["active"]["type"] == "BOOL" + + +def test_format_search_describe_response_with_all_options() -> None: + # Test describe response with all field options + raw_response = [ + "name", "fullindex", + "type", "HASH", + "schema", [ + ["title", "TEXT", "NOSTEM", "NOTOKENIZE"], + ["score", "I64", "FAST"], + ], + ] + result = format_search_describe_response(raw_response, None) + + assert result["name"] == "fullindex" + assert result["dataType"] == "hash" + assert result["schema"]["title"]["type"] == "TEXT" + assert result["schema"]["title"]["noStem"] is True + assert result["schema"]["title"]["noTokenize"] is True + assert result["schema"]["score"]["type"] == "I64" + assert result["schema"]["score"]["fast"] is True + + +def test_format_search_count_response() -> None: + # Test count response with int + assert format_search_count_response(42, None) == {"count": 42} + + # Test count response with string + assert format_search_count_response("100", None) == {"count": 100} diff --git a/upstash_redis/asyncio/client.py b/upstash_redis/asyncio/client.py index 2fcfaf1..5de8cda 100644 --- a/upstash_redis/asyncio/client.py +++ b/upstash_redis/asyncio/client.py @@ -6,11 +6,11 @@ PipelineCommands, AsyncJsonCommands, PipelineJsonCommands, + AsyncSearchCommands, ) from upstash_redis.format import cast_response from upstash_redis.http import make_headers, AsyncHttpClient from upstash_redis.typing import RESTResultT -from upstash_redis.search_namespace import SearchNamespace class Redis(AsyncCommands): @@ -61,7 +61,7 @@ def __init__( self._headers = make_headers(token, rest_encoding, allow_telemetry) self._json = AsyncJsonCommands(self) - self._search = SearchNamespace(self) + self._search = AsyncSearchCommands(self) self._http = AsyncHttpClient( encoding=rest_encoding, retries=rest_retries, @@ -74,7 +74,7 @@ def json(self) -> AsyncJsonCommands: return self._json @property - def search(self) -> SearchNamespace: + def search(self) -> AsyncSearchCommands: return self._search @classmethod diff --git a/upstash_redis/client.py b/upstash_redis/client.py index 76b473f..53c0fac 100644 --- a/upstash_redis/client.py +++ b/upstash_redis/client.py @@ -6,11 +6,11 @@ PipelineCommands, JsonCommands, PipelineJsonCommands, + SearchCommands, ) from upstash_redis.format import cast_response from upstash_redis.http import make_headers, SyncHttpClient from upstash_redis.typing import RESTResultT -from upstash_redis.search_namespace import SearchNamespace class Redis(Commands): @@ -63,7 +63,7 @@ def __init__( self._headers = make_headers(token, rest_encoding, allow_telemetry) self._json = JsonCommands(self) - self._search = SearchNamespace(self) + self._search = SearchCommands(self) self._http = SyncHttpClient( encoding=rest_encoding, retries=rest_retries, @@ -76,7 +76,7 @@ def json(self) -> JsonCommands: return self._json @property - def search(self) -> SearchNamespace: + def search(self) -> SearchCommands: return self._search @classmethod diff --git a/upstash_redis/commands.py b/upstash_redis/commands.py index 34fa5b5..aaba0c0 100644 --- a/upstash_redis/commands.py +++ b/upstash_redis/commands.py @@ -9,6 +9,8 @@ handle_zrangebylex_exceptions, number_are_not_none, ) +from upstash_redis.search_types import CreateIndexParams, QueryOptions +from upstash_redis.search_utils import build_create_index_command, build_query_command ResponseT = Union[Awaitable, Any] @@ -5484,6 +5486,153 @@ def type(self, key: str, path: str = "$") -> ResponseT: return self.client.execute(command=command) +class SearchCommands: + """ + Search commands namespace. + + This class provides methods for search index operations. + It follows the same pattern as JsonCommands. + """ + def __init__(self, client: Commands): + self.client = client + + def create_index( + self, + *, + name: str, + schema, + dataType: str, + prefix, + language = None, + skipInitialScan: bool = False, + existsOk: bool = False, + ) -> "SearchIndexCommands": + """Create a new search index.""" + params: CreateIndexParams = { + "name": name, + "schema": schema, + "dataType": dataType, + "prefix": prefix, + } + if language is not None: + params["language"] = language + if skipInitialScan: + params["skipInitialScan"] = skipInitialScan + if existsOk: + params["existsOk"] = existsOk + + command = build_create_index_command(params) + result = self.client.execute(command) + + # Handle async + if hasattr(result, "__await__"): + async def _async_create(): + await result + return SearchIndexCommands(self.client, name, schema) + return _async_create() + + return SearchIndexCommands(self.client, name, schema) + + def index(self, name: str, schema = None) -> "SearchIndexCommands": + """Initialize a SearchIndexCommands instance for an existing index.""" + return SearchIndexCommands(self.client, name, schema) + + +class SearchIndexCommands: + """ + Commands for interacting with a specific search index. + + This class provides methods to query, count, describe, and drop an index. + """ + def __init__(self, client: Commands, name: str, schema = None): + self.client = client + self.name = name + self.schema = schema + + def wait_indexing(self) -> ResponseT: + """ + Wait for the index to finish indexing all existing keys. + + Example: + ```python + index.wait_indexing() + ``` + """ + command = ["SEARCH.WAITINDEXING", self.name] + return self.client.execute(command) + + def describe(self) -> ResponseT: + """ + Get detailed information about the index structure. + + Example: + ```python + description = index.describe() + ``` + """ + command = ["SEARCH.DESCRIBE", self.name] + return self.client.execute(command) + + def query( + self, + *, + filter, + limit = None, + offset = None, + orderBy = None, + select = None, + highlight = None, + ) -> ResponseT: + """ + Query the index with filters and options. + + Example: + ```python + results = index.query(filter={"name": {"$eq": "Laptop"}}) + ``` + """ + options: QueryOptions = {} + if filter is not None: + options["filter"] = filter + if limit is not None: + options["limit"] = limit + if offset is not None: + options["offset"] = offset + if orderBy is not None: + options["orderBy"] = orderBy + if select is not None: + options["select"] = select + if highlight is not None: + options["highlight"] = highlight + + command = build_query_command("SEARCH.QUERY", self.name, options) + return self.client.execute(command) + + def count(self, filter) -> ResponseT: + """ + Count documents matching a filter. + + Example: + ```python + result = index.count({"active": {"$eq": True}}) + ``` + """ + command = build_query_command("SEARCH.COUNT", self.name, {"filter": filter}) + return self.client.execute(command) + + def drop(self) -> ResponseT: + """ + Drop (delete) the index. + + Example: + ```python + result = index.drop() + ``` + """ + command = ["SEARCH.DROP", self.name] + return self.client.execute(command) + + # It doesn't inherit from "Redis" mainly because of the methods signatures. class BitFieldCommands: def __init__(self, client: Commands, key: str): @@ -5574,6 +5723,8 @@ def execute(self) -> ResponseT: AsyncCommands = Commands AsyncJsonCommands = JsonCommands +AsyncSearchCommands = SearchCommands +AsyncSearchIndexCommands = SearchIndexCommands AsyncBitFieldCommands = BitFieldCommands AsyncBitFieldROCommands = BitFieldROCommands PipelineCommands = Commands diff --git a/upstash_redis/commands.pyi b/upstash_redis/commands.pyi index 5b3b60c..3f06f74 100644 --- a/upstash_redis/commands.pyi +++ b/upstash_redis/commands.pyi @@ -3,6 +3,7 @@ from typing import Any, Dict, List, Literal, Mapping, Optional, Tuple, Union from upstash_redis.typing import FloatMinMaxT, ValueT, JSONValueT from upstash_redis.utils import GeoSearchResult +from upstash_redis.search_types import NestedIndexSchema, FlatIndexSchema class Commands: def bitcount( @@ -2028,3 +2029,75 @@ class AsyncJsonCommands: async def strlen(self, key: str, path: str = "$") -> List[Union[int, None]]: ... async def toggle(self, key: str, path: str = "$") -> List[Union[bool, None]]: ... async def type(self, key: str, path: str = "$") -> List[str]: ... + +class SearchCommands: + def __init__(self, client: Commands): ... + def create_index( + self, + *, + name: str, + schema: Union[NestedIndexSchema, FlatIndexSchema], + dataType: str, + prefix: Union[str, List[str]], + language: Optional[str] = None, + skipInitialScan: bool = False, + existsOk: bool = False, + ) -> SearchIndexCommands: ... + def index( + self, + name: str, + schema: Optional[Union[NestedIndexSchema, FlatIndexSchema]] = None + ) -> SearchIndexCommands: ... + +class SearchIndexCommands: + def __init__(self, client: Commands, name: str, schema: Optional[Union[NestedIndexSchema, FlatIndexSchema]] = None): ... + def wait_indexing(self) -> Any: ... + def describe(self) -> Dict[str, Any]: ... + def query( + self, + *, + filter: Dict[str, Any], + limit: Optional[int] = None, + offset: Optional[int] = None, + orderBy: Optional[Dict[str, str]] = None, + select: Optional[Dict[str, bool]] = None, + highlight: Optional[Dict[str, Any]] = None, + ) -> List[Dict[str, Any]]: ... + def count(self, filter: Dict[str, Any]) -> Dict[str, int]: ... + def drop(self) -> int: ... + +class AsyncSearchCommands: + def __init__(self, client: AsyncCommands): ... + async def create_index( + self, + *, + name: str, + schema: Union[NestedIndexSchema, FlatIndexSchema], + dataType: str, + prefix: Union[str, List[str]], + language: Optional[str] = None, + skipInitialScan: bool = False, + existsOk: bool = False, + ) -> AsyncSearchIndexCommands: ... + def index( + self, + name: str, + schema: Optional[Union[NestedIndexSchema, FlatIndexSchema]] = None + ) -> AsyncSearchIndexCommands: ... + +class AsyncSearchIndexCommands: + def __init__(self, client: AsyncCommands, name: str, schema: Optional[Union[NestedIndexSchema, FlatIndexSchema]] = None): ... + async def wait_indexing(self) -> Any: ... + async def describe(self) -> Dict[str, Any]: ... + async def query( + self, + *, + filter: Dict[str, Any], + limit: Optional[int] = None, + offset: Optional[int] = None, + orderBy: Optional[Dict[str, str]] = None, + select: Optional[Dict[str, bool]] = None, + highlight: Optional[Dict[str, Any]] = None, + ) -> List[Dict[str, Any]]: ... + async def count(self, filter: Dict[str, Any]) -> Dict[str, int]: ... + async def drop(self) -> int: ... diff --git a/upstash_redis/format.py b/upstash_redis/format.py index 0e1e6d1..4bf87df 100644 --- a/upstash_redis/format.py +++ b/upstash_redis/format.py @@ -1,5 +1,5 @@ from json import loads -from typing import Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, List, Optional, Tuple, Union from upstash_redis.typing import RESTResultT from upstash_redis.utils import GeoSearchResult @@ -226,6 +226,110 @@ def format_xpending_response(res, command): return res +def format_search_query_response(res: List[Any], _) -> List[Dict[str, Any]]: + """Format SEARCH.QUERY response into structured results.""" + results: List[Dict[str, Any]] = [] + + for item_raw in res: + if not isinstance(item_raw, (list, tuple)) or len(item_raw) < 2: + continue + + key = item_raw[0] + score = item_raw[1] + + result: Dict[str, Any] = {"key": key, "score": score} + + if len(item_raw) > 2: + raw_fields = item_raw[2] + + if isinstance(raw_fields, (list, tuple)) and len(raw_fields) > 0: + data: Dict[str, Any] = {} + + # Process field pairs + for field_raw in raw_fields: + if not isinstance(field_raw, (list, tuple)) or len(field_raw) < 2: + continue + + field_key = field_raw[0] + field_value = field_raw[1] + + # Handle nested paths + path_parts = field_key.split(".") + if len(path_parts) == 1: + data[field_key] = field_value + else: + # Build nested structure + current_obj = data + for i, part in enumerate(path_parts[:-1]): + if part not in current_obj: + current_obj[part] = {} + current_obj = current_obj[part] + current_obj[path_parts[-1]] = field_value + + # If $ key exists (full document), use its contents + if "$" in data: + data = data["$"] + + result["data"] = data + + results.append(result) + + return results + + +def format_search_describe_response(res: List[Any], _) -> Dict[str, Any]: + """Format SEARCH.DESCRIBE response into index description.""" + description: Dict[str, Any] = {} + + i = 0 + while i < len(res): + descriptor = res[i] + value = res[i + 1] if i + 1 < len(res) else None + + if descriptor == "name": + description["name"] = value + elif descriptor == "type": + if value: + description["dataType"] = value.lower() + elif descriptor == "prefixes": + description["prefixes"] = value if isinstance(value, list) else [] + elif descriptor == "language": + description["language"] = value + elif descriptor == "schema": + schema: Dict[str, Dict[str, Any]] = {} + if isinstance(value, list): + for field_desc in value: + if not isinstance(field_desc, list) or len(field_desc) < 2: + continue + + field_name = field_desc[0] + field_info: Dict[str, Any] = {"type": field_desc[1]} + + # Parse field options + for j in range(2, len(field_desc)): + option = field_desc[j] + if option == "NOSTEM": + field_info["noStem"] = True + elif option == "NOTOKENIZE": + field_info["noTokenize"] = True + elif option == "FAST": + field_info["fast"] = True + + schema[field_name] = field_info + + description["schema"] = schema + + i += 2 + + return description + + +def format_search_count_response(res: Any, _) -> Dict[str, int]: + """Format SEARCH.COUNT response.""" + count = int(res) if not isinstance(res, int) else res + return {"count": count} + + FORMATTERS: Dict[str, Callable] = { "COPY": to_bool, "EXPIRE": to_bool, @@ -301,6 +405,9 @@ def format_xpending_response(res, command): "XPENDING": format_xpending_response, "XGROUP DESTROY": to_bool, "XGROUP CREATECONSUMER": to_bool, + "SEARCH.QUERY": format_search_query_response, + "SEARCH.DESCRIBE": format_search_describe_response, + "SEARCH.COUNT": format_search_count_response, } diff --git a/upstash_redis/search_index.py b/upstash_redis/search_index.py deleted file mode 100644 index 6a86449..0000000 --- a/upstash_redis/search_index.py +++ /dev/null @@ -1,281 +0,0 @@ -"""Search index functionality for Upstash Redis.""" - -from typing import Any, Dict, List, Optional, Union - -from upstash_redis.search_types import ( - CreateIndexParams, - IndexDescription, - QueryOptions, - QueryResult, - NestedIndexSchema, - FlatIndexSchema, - RootQueryFilter, -) -from upstash_redis.search_utils import ( - build_create_index_command, - build_query_command, - deserialize_describe_response, - deserialize_query_response, - parse_count_response, -) - - -class SearchIndex: - """ - Represents a search index in Upstash Redis. - - This class provides methods to interact with a search index including - querying, counting, describing, and dropping the index. - """ - - def __init__( - self, - name: str, - client: Any, - schema: Optional[Union[NestedIndexSchema, FlatIndexSchema]] = None, - ): - """ - Initialize a SearchIndex instance. - - Args: - name: Name of the index - client: Redis client instance (sync or async) - schema: Optional schema definition for the index - """ - self.name = name - self.schema = schema - self._client = client - - def wait_indexing(self) -> Any: - """ - Wait for the index to finish indexing all existing keys. - - Returns: - The result of the command execution - - Example: - ```python - index.wait_indexing() - ``` - """ - command = ["SEARCH.WAITINDEXING", self.name] - return self._client.execute(command) - - def describe(self) -> IndexDescription: - """ - Get detailed information about the index structure. - - Returns: - IndexDescription with schema, prefixes, and other metadata - - Example: - ```python - description = index.describe() - print(description["schema"]) - ``` - """ - command = ["SEARCH.DESCRIBE", self.name] - raw_result = self._client.execute(command) - - # If the client returns a coroutine (async), we can't deserialize here - # The async version will need to override this method - if hasattr(raw_result, "__await__"): - # For async client, return the coroutine wrapped in deserializer - async def _async_describe(): - result = await raw_result - return deserialize_describe_response(result) - - return _async_describe() - - return deserialize_describe_response(raw_result) - - def query( - self, - *, - filter: RootQueryFilter, - limit: Optional[int] = None, - offset: Optional[int] = None, - orderBy: Optional[Dict[str, str]] = None, - select: Optional[Dict[str, bool]] = None, - highlight: Optional[Dict[str, Any]] = None, - ) -> List[QueryResult]: - """ - Query the index with filters and options. - - Args: - filter: Filter specification mapping field names to conditions - limit: Maximum number of results to return - offset: Number of results to skip (for pagination) - orderBy: Field to sort by with direction ("ASC" or "DESC") - select: Fields to include/exclude in results - highlight: Highlighting options for matched terms - - Returns: - List of query results with keys, scores, and data - - Example: - ```python - # Simple text search - results = index.query(filter={"name": {"$eq": "Laptop"}}) - - # With pagination and sorting - results = index.query( - filter={"category": {"$eq": "electronics"}}, - limit=10, - offset=0, - orderBy={"price": "ASC"} - ) - ``` - """ - options: QueryOptions = {"filter": filter} - if limit is not None: - options["limit"] = limit - if offset is not None: - options["offset"] = offset - if orderBy is not None: - options["orderBy"] = orderBy - if select is not None: - options["select"] = select - if highlight is not None: - options["highlight"] = highlight - - command = build_query_command("SEARCH.QUERY", self.name, options) - raw_result = self._client.execute(command) - - # Handle async - if hasattr(raw_result, "__await__"): - - async def _async_query(): - result = await raw_result - return deserialize_query_response(result) - - return _async_query() - - return deserialize_query_response(raw_result) - - def count(self, filter: RootQueryFilter) -> Dict[str, int]: - """ - Count documents matching a filter. - - Args: - filter: Filter specification mapping field names to conditions - - Returns: - Dictionary with "count" key containing the number of matches - - Example: - ```python - result = index.count({"active": {"$eq": True}}) - print(result["count"]) # e.g., 42 - ``` - """ - command = build_query_command("SEARCH.COUNT", self.name, {"filter": filter}) - raw_result = self._client.execute(command) - - # Handle async - if hasattr(raw_result, "__await__"): - - async def _async_count(): - result = await raw_result - return {"count": parse_count_response(result)} - - return _async_count() - - return {"count": parse_count_response(raw_result)} - - def drop(self) -> Union[int, Any]: - """ - Drop (delete) the index. - - Returns: - 1 if the index was dropped, 0 if it didn't exist - - Example: - ```python - result = index.drop() - print(result) # 1 - ``` - """ - command = ["SEARCH.DROP", self.name] - return self._client.execute(command) - - -def _create_index(client: Any, params: CreateIndexParams) -> SearchIndex: - """ - Create a new search index. - - Args: - client: Redis client instance (sync or async) - params: Index creation parameters including name, schema, data type, etc. - - Returns: - SearchIndex instance for the newly created index - - Example: - ```python - schema = { - "name": "TEXT", - "age": {"type": "U64", "fast": True}, - "active": "BOOL" - } - - index = create_index(redis, { - "name": "users-idx", - "prefix": "user:", - "dataType": "json", - "schema": schema - }) - ``` - """ - name = params["name"] - schema = params.get("schema") - - # Build and execute create command - command = build_create_index_command(params) - result = client.execute(command) - - # Handle async - if hasattr(result, "__await__"): - - async def _async_create(): - await result - return SearchIndex(name=name, client=client, schema=schema) - - return _async_create() - - return SearchIndex(name=name, client=client, schema=schema) - - -def _init_index( - client: Any, - name: str, - schema: Optional[Union[NestedIndexSchema, FlatIndexSchema]] = None, -) -> SearchIndex: - """ - Initialize a SearchIndex instance for an existing index. - - This function does not create the index, it only creates a Python object - to interact with an existing index. - - Args: - client: Redis client instance (sync or async) - name: Name of the existing index - schema: Optional schema definition for better type hints - - Returns: - SearchIndex instance - - Example: - ```python - # Without schema - index = init_index(redis, "users-idx") - - # With schema for better typing - schema = { - "name": "TEXT", - "age": {"type": "U64", "fast": True} - } - index = init_index(redis, "users-idx", schema) - ``` - """ - return SearchIndex(name=name, client=client, schema=schema) diff --git a/upstash_redis/search_namespace.py b/upstash_redis/search_namespace.py deleted file mode 100644 index f89c61b..0000000 --- a/upstash_redis/search_namespace.py +++ /dev/null @@ -1,113 +0,0 @@ -"""Search namespace for Redis client.""" - -from typing import Any, Optional, Union - -from upstash_redis.search_index import SearchIndex, _create_index, _init_index -from upstash_redis.search_types import CreateIndexParams, NestedIndexSchema, FlatIndexSchema - - -class SearchNamespace: - """ - Namespace for search index operations. - - Access via redis.search.create_index() or redis.search.index() - """ - - def __init__(self, client: Any): - """ - Initialize the search namespace. - - Args: - client: Redis client instance (sync or async) - """ - self._client = client - - def create_index( - self, - *, - name: str, - schema: Union[NestedIndexSchema, FlatIndexSchema], - dataType: str, - prefix: Union[str, list[str]], - language: Optional[str] = None, - skipInitialScan: bool = False, - existsOk: bool = False, - ) -> SearchIndex: - """ - Create a new search index. - - Args: - name: Name of the index - schema: Schema definition mapping field names to types - dataType: Type of data being indexed ("string", "json", or "hash") - prefix: Key prefix(es) for automatic indexing - language: Optional language for text analysis - skipInitialScan: Skip indexing existing keys - existsOk: Don't error if index already exists - - Returns: - SearchIndex instance for the newly created index - - Example: - ```python - schema = { - "name": "TEXT", - "age": {"type": "U64", "fast": True}, - "active": "BOOL" - } - - index = redis.search.create_index( - name="users-idx", - prefix="user:", - dataType="json", - schema=schema - ) - ``` - """ - params: CreateIndexParams = { - "name": name, - "schema": schema, - "dataType": dataType, - "prefix": prefix, - } - if language is not None: - params["language"] = language - if skipInitialScan: - params["skipInitialScan"] = skipInitialScan - if existsOk: - params["existsOk"] = existsOk - - return _create_index(self._client, params) - - def index( - self, - name: str, - schema: Optional[Union[NestedIndexSchema, FlatIndexSchema]] = None - ) -> SearchIndex: - """ - Initialize a SearchIndex instance for an existing index. - - This method does not create the index, it only creates a Python object - to interact with an existing index. - - Args: - name: Name of the existing index - schema: Optional schema definition for better type hints - - Returns: - SearchIndex instance - - Example: - ```python - # Without schema - index = redis.search.index("users-idx") - - # With schema for better typing - schema = { - "name": "TEXT", - "age": {"type": "U64", "fast": True} - } - index = redis.search.index("users-idx", schema) - ``` - """ - return _init_index(self._client, name, schema) From d86e5d46ff036dce15cb0a7b4310e4b55516f6ed Mon Sep 17 00:00:00 2001 From: CahidArda Date: Wed, 21 Jan 2026 16:31:42 +0300 Subject: [PATCH 04/12] fix: add formatter for create_index --- .../asyncio/search/test_search_index_async.py | 2 +- tests/commands/search/test_search_index.py | 13 +-- tests/test_formatters.py | 30 +++---- upstash_redis/asyncio/client.py | 8 +- upstash_redis/client.py | 8 +- upstash_redis/commands.py | 20 ++--- upstash_redis/format.py | 84 +++++++++++-------- 7 files changed, 83 insertions(+), 82 deletions(-) diff --git a/tests/commands/asyncio/search/test_search_index_async.py b/tests/commands/asyncio/search/test_search_index_async.py index 39713d7..c8e6146 100644 --- a/tests/commands/asyncio/search/test_search_index_async.py +++ b/tests/commands/asyncio/search/test_search_index_async.py @@ -44,7 +44,7 @@ class TestAsyncQuery: """Tests for async query operations.""" @pytest_asyncio.fixture - async def query_index(self, async_redis_client, cleanup_indexes): + async def query_index(self, async_redis_client: Redis, cleanup_indexes): """Create a test index with data for querying.""" name = f"test-async-query-{random_id()}" prefix = f"{name}:" diff --git a/tests/commands/search/test_search_index.py b/tests/commands/search/test_search_index.py index e8150d1..9abe34e 100644 --- a/tests/commands/search/test_search_index.py +++ b/tests/commands/search/test_search_index.py @@ -478,20 +478,9 @@ def test_drop_existing_index(self, redis_client: Redis): class TestIndexMethod: """Tests for redis.search.index method.""" - def test_index_without_schema(self, redis_client: Redis): + def test_index(self, redis_client: Redis): """Test creating a SearchIndex instance without schema.""" index = redis_client.search.index("test-index") assert index is not None assert index.name == "test-index" - assert index.schema is None - - def test_index_with_schema(self, redis_client: Redis): - """Test creating a SearchIndex instance with schema.""" - schema = {"name": "TEXT", "age": {"type": "U64", "fast": True}} - - index = redis_client.search.index("test-index", schema) - - assert index is not None - assert index.name == "test-index" - assert index.schema == schema diff --git a/tests/test_formatters.py b/tests/test_formatters.py index 32b0bb9..c6fb18b 100644 --- a/tests/test_formatters.py +++ b/tests/test_formatters.py @@ -15,7 +15,7 @@ def test_list_to_dict() -> None: - assert to_dict(["a", "1", "b", "2", "c", 3], None) == { + assert to_dict(["a", "1", "b", "2", "c", 3], None, None) == { "a": "1", "b": "2", "c": 3, @@ -23,15 +23,15 @@ def test_list_to_dict() -> None: def test_format_float_list() -> None: - assert to_optional_float_list(["1.1", "2.2", None], None) == [1.1, 2.2, None] + assert to_optional_float_list(["1.1", "2.2", None], None, None) == [1.1, 2.2, None] def test_string_to_json() -> None: - assert string_to_json('{"a": 1, "b": 2}', None) == {"a": 1, "b": 2} + assert string_to_json('{"a": 1, "b": 2}', None, None) == {"a": 1, "b": 2} def test_string_list_to_json_list() -> None: - assert to_json_list(['{"a": 1, "b": 2}', '{"c": 3, "d": 4}'], None) == [ + assert to_json_list(['{"a": 1, "b": 2}', '{"c": 3, "d": 4}'], None, None) == [ {"a": 1, "b": 2}, {"c": 3, "d": 4}, ] @@ -155,7 +155,7 @@ def test_format_geo_members_with_hash_and_coordinates() -> None: def test_format_geo_positions() -> None: - assert format_geopos([["1.0", "2.5"], ["3.1", "4.2"], None], None) == [ + assert format_geopos([["1.0", "2.5"], ["3.1", "4.2"], None], None, None) == [ (1.0, 2.5), (3.1, 4.2), None, @@ -163,11 +163,11 @@ def test_format_geo_positions() -> None: def test_format_server_time() -> None: - assert format_time(["1620752099", "12"], None) == (1620752099, 12) + assert format_time(["1620752099", "12"], None, None) == (1620752099, 12) def test_list_to_optional_bool_list() -> None: - assert to_optional_bool_list([1, 0, None, 1], None) == [ + assert to_optional_bool_list([1, 0, None, 1], None, None) == [ True, False, None, @@ -181,7 +181,7 @@ def test_format_search_query_response() -> None: ["doc:1", 0.95, [["name", "Laptop"], ["price", "999"]]], ["doc:2", 0.85, [["name", "Phone"], ["price", "699"]]], ] - result = format_search_query_response(raw_response, None) + result = format_search_query_response(raw_response, None, None) assert len(result) == 2 assert result[0]["key"] == "doc:1" @@ -199,7 +199,7 @@ def test_format_search_query_response_with_nested_paths() -> None: raw_response = [ ["doc:1", 0.95, [["user.name", "John"], ["user.age", "30"]]], ] - result = format_search_query_response(raw_response, None) + result = format_search_query_response(raw_response, None, None) assert len(result) == 1 assert result[0]["key"] == "doc:1" @@ -213,7 +213,7 @@ def test_format_search_query_response_with_dollar_key() -> None: raw_response = [ ["doc:1", 1.0, [["$", {"name": "Laptop", "price": 999}]]], ] - result = format_search_query_response(raw_response, None) + result = format_search_query_response(raw_response, None, None) assert len(result) == 1 assert result[0]["key"] == "doc:1" @@ -228,7 +228,7 @@ def test_format_search_query_response_without_fields() -> None: ["doc:1", 0.95], ["doc:2", 0.85], ] - result = format_search_query_response(raw_response, None) + result = format_search_query_response(raw_response, None, None) assert len(result) == 2 assert result[0]["key"] == "doc:1" @@ -252,7 +252,7 @@ def test_format_search_describe_response() -> None: ["active", "BOOL"], ], ] - result = format_search_describe_response(raw_response, None) + result = format_search_describe_response(raw_response, None, None) assert result["name"] == "myindex" assert result["dataType"] == "json" @@ -278,7 +278,7 @@ def test_format_search_describe_response_with_all_options() -> None: ["score", "I64", "FAST"], ], ] - result = format_search_describe_response(raw_response, None) + result = format_search_describe_response(raw_response, None, None) assert result["name"] == "fullindex" assert result["dataType"] == "hash" @@ -291,7 +291,7 @@ def test_format_search_describe_response_with_all_options() -> None: def test_format_search_count_response() -> None: # Test count response with int - assert format_search_count_response(42, None) == {"count": 42} + assert format_search_count_response(42, None, None) == {"count": 42} # Test count response with string - assert format_search_count_response("100", None) == {"count": 100} + assert format_search_count_response("100", None, None) == {"count": 100} diff --git a/upstash_redis/asyncio/client.py b/upstash_redis/asyncio/client.py index 5de8cda..2af7bd1 100644 --- a/upstash_redis/asyncio/client.py +++ b/upstash_redis/asyncio/client.py @@ -142,7 +142,7 @@ async def execute(self, command: List) -> RESTResultT: headers=self._headers, command=command, ) - return cast_response(command, res) + return cast_response(command, res, self) def pipeline(self) -> "AsyncPipeline": """ @@ -154,6 +154,7 @@ def pipeline(self) -> "AsyncPipeline": http=self._http, multi_exec="pipeline", set_sync_token_header_fn=self._maybe_set_sync_token_header, + original_client=self ) def multi(self) -> "AsyncPipeline": @@ -166,6 +167,7 @@ def multi(self) -> "AsyncPipeline": http=self._http, multi_exec="multi-exec", set_sync_token_header_fn=self._maybe_set_sync_token_header, + original_client=self ) @@ -177,6 +179,7 @@ def __init__( http: AsyncHttpClient, multi_exec: Literal["multi-exec", "pipeline"], set_sync_token_header_fn: Callable[[Dict[str, str]], None], + original_client: Optional[Redis], ): self._url = f"{url}/{multi_exec}" self._headers = headers @@ -186,6 +189,7 @@ def __init__( self._command_stack: List[List[str]] = [] self._set_sync_token_header_fn = set_sync_token_header_fn + self._original_client = original_client @property def json(self) -> PipelineJsonCommands: @@ -214,7 +218,7 @@ async def exec(self) -> List[RESTResultT]: from_pipeline=True, ) response = [ - cast_response(command, response) + cast_response(command, response, self._original_client) for command, response in zip(self._command_stack, res) ] self.reset() diff --git a/upstash_redis/client.py b/upstash_redis/client.py index 53c0fac..ac88d64 100644 --- a/upstash_redis/client.py +++ b/upstash_redis/client.py @@ -144,7 +144,7 @@ def execute(self, command: List) -> RESTResultT: headers=self._headers, command=command, ) - return cast_response(command, res) + return cast_response(command, res, self) def pipeline(self) -> "Pipeline": """ @@ -156,6 +156,7 @@ def pipeline(self) -> "Pipeline": http=self._http, multi_exec="pipeline", set_sync_token_header_fn=self._maybe_set_sync_token_header, + original_client=self ) def multi(self) -> "Pipeline": @@ -168,6 +169,7 @@ def multi(self) -> "Pipeline": http=self._http, multi_exec="multi-exec", set_sync_token_header_fn=self._maybe_set_sync_token_header, + original_client=self ) @@ -179,6 +181,7 @@ def __init__( http: SyncHttpClient, multi_exec: Literal["multi-exec", "pipeline"], set_sync_token_header_fn: Callable[[Dict[str, str]], None], + original_client: Optional[Redis], ): self._url = f"{url}/{multi_exec}" self._headers = headers @@ -188,6 +191,7 @@ def __init__( self._command_stack: List[List[str]] = [] self._set_sync_token_header_fn = set_sync_token_header_fn + self._original_client = original_client @property def json(self) -> PipelineJsonCommands: @@ -216,7 +220,7 @@ def exec(self) -> List[RESTResultT]: from_pipeline=True, ) response = [ - cast_response(command, response) + cast_response(command, response, self._original_client) for command, response in zip(self._command_stack, res) ] self.reset() diff --git a/upstash_redis/commands.py b/upstash_redis/commands.py index aaba0c0..b2d3ac8 100644 --- a/upstash_redis/commands.py +++ b/upstash_redis/commands.py @@ -5506,7 +5506,7 @@ def create_index( language = None, skipInitialScan: bool = False, existsOk: bool = False, - ) -> "SearchIndexCommands": + ) -> ResponseT: """Create a new search index.""" params: CreateIndexParams = { "name": name, @@ -5522,20 +5522,11 @@ def create_index( params["existsOk"] = existsOk command = build_create_index_command(params) - result = self.client.execute(command) - - # Handle async - if hasattr(result, "__await__"): - async def _async_create(): - await result - return SearchIndexCommands(self.client, name, schema) - return _async_create() - - return SearchIndexCommands(self.client, name, schema) + return self.client.execute(command) - def index(self, name: str, schema = None) -> "SearchIndexCommands": + def index(self, name: str) -> "SearchIndexCommands": """Initialize a SearchIndexCommands instance for an existing index.""" - return SearchIndexCommands(self.client, name, schema) + return SearchIndexCommands(self.client, name) class SearchIndexCommands: @@ -5544,10 +5535,9 @@ class SearchIndexCommands: This class provides methods to query, count, describe, and drop an index. """ - def __init__(self, client: Commands, name: str, schema = None): + def __init__(self, client: Commands, name: str): self.client = client self.name = name - self.schema = schema def wait_indexing(self) -> ResponseT: """ diff --git a/upstash_redis/format.py b/upstash_redis/format.py index 4bf87df..83f21a3 100644 --- a/upstash_redis/format.py +++ b/upstash_redis/format.py @@ -3,9 +3,10 @@ from upstash_redis.typing import RESTResultT from upstash_redis.utils import GeoSearchResult +from upstash_redis.commands import SearchIndexCommands -def to_dict(res: List, _) -> Dict: +def to_dict(res: List, _, __) -> Dict: """ Convert a list that contains ungrouped pairs as consecutive elements (usually field-value or similar) into a dict. """ @@ -14,7 +15,7 @@ def to_dict(res: List, _) -> Dict: def format_geopos( - res: List[Optional[List[str]]], _ + res: List[Optional[List[str]]], _, __ ) -> List[Union[Tuple[float, float], None]]: return [ (float(member[0]), float(member[1])) if isinstance(member, List) else None @@ -83,11 +84,11 @@ def format_geo_search_response( return results -def format_time(res: List[str], _) -> Tuple[int, int]: +def format_time(res: List[str], _, __) -> Tuple[int, int]: return int(res[0]), int(res[1]) -def format_sorted_set_response(res: List[str], _) -> List[Tuple[str, float]]: +def format_sorted_set_response(res: List[str], _, __) -> List[Tuple[str, float]]: """ Format the raw output given by Sorted Set commands, usually the ones that return the member-score pairs of Sorted Sets. @@ -96,7 +97,7 @@ def format_sorted_set_response(res: List[str], _) -> List[Tuple[str, float]]: return list(zip(it, map(float, it))) -def to_optional_float_list(res: List[Optional[str]], _) -> List[Optional[float]]: +def to_optional_float_list(res: List[Optional[str]], _, __) -> List[Optional[float]]: """ Format a list of strings representing floats or None values. """ @@ -104,7 +105,7 @@ def to_optional_float_list(res: List[Optional[str]], _) -> List[Optional[float]] return [float(value) if value is not None else None for value in res] -def format_set(res, command): +def format_set(res, command, _): options = command[3:] if "GET" in options: @@ -112,61 +113,61 @@ def format_set(res, command): return res == "OK" -def string_to_json(res, _): +def string_to_json(res, _, __): if res is None: return None return loads(res) -def to_json_list(res, command): - return [string_to_json(value, command) for value in res] +def to_json_list(res, command, client): + return [string_to_json(value, command, client) for value in res] -def ok_to_bool(res, _): +def ok_to_bool(res, _, __): return res == "OK" -def to_bool(res, _): +def to_bool(res, _, __): return bool(res) -def to_bool_list(res, _): +def to_bool_list(res, _, __): return list(map(bool, res)) -def to_optional_bool_list(res, _): +def to_optional_bool_list(res, _, __): return [bool(value) if value is not None else None for value in res] -def to_optional_float(res, _): +def to_optional_float(res, _, __): if res is None: return None return float(res) -def to_float(res, _): +def to_float(res, _, __): return float(res) -def format_scan(res, _): +def format_scan(res, _, __): return int(res[0]), res[1] -def format_hscan(res, _): - return int(res[0]), to_dict(res[1], None) +def format_hscan(res, _, client): + return int(res[0]), to_dict(res[1], None, client) -def format_zscan(res, _): - return int(res[0]), format_sorted_set_response(res[1], None) +def format_zscan(res, _, client): + return int(res[0]), format_sorted_set_response(res[1], None, client) -def format_zscore(res, _): +def format_zscore(res, _, __): return float(res) if res is not None else res -def format_search(res, command): +def format_search(res, command, _): withdist = "WITHDIST" in command withhash = "WITHHASH" in command withcoord = "WITHCOORD" in command @@ -176,15 +177,15 @@ def format_search(res, command): return res -def format_hrandfield(res, command): +def format_hrandfield(res, command, client): with_values = "WITHVALUES" in command if with_values: - return to_dict(res, command) + return to_dict(res, command, client) return res -def format_zadd(res, command): +def format_zadd(res, command, _): incr = "INCR" in command if incr: return float(res) if res is not None else res @@ -192,15 +193,15 @@ def format_zadd(res, command): return res -def format_sorted_set_response_with_score(res, command): +def format_sorted_set_response_with_score(res, command, client): with_scores = "WITHSCORES" in command if with_scores: - return format_sorted_set_response(res, command) + return format_sorted_set_response(res, command, client) return res -def format_xread_response(res, command): +def format_xread_response(res, command, _): """Format XREAD/XREADGROUP response to convert stream entries.""" if res is None: return [] @@ -211,7 +212,7 @@ def format_xread_response(res, command): return res -def format_xrange_response(res, command): +def format_xrange_response(res, command, _): """Format XRANGE/XREVRANGE response to convert entries.""" if not res: return res @@ -220,13 +221,13 @@ def format_xrange_response(res, command): return res -def format_xpending_response(res, command): +def format_xpending_response(res, command, _): """Format XPENDING response.""" # Return as-is to maintain Redis-compatible format return res -def format_search_query_response(res: List[Any], _) -> List[Dict[str, Any]]: +def format_search_query_response(res: List[Any], _, __) -> List[Dict[str, Any]]: """Format SEARCH.QUERY response into structured results.""" results: List[Dict[str, Any]] = [] @@ -277,7 +278,7 @@ def format_search_query_response(res: List[Any], _) -> List[Dict[str, Any]]: return results -def format_search_describe_response(res: List[Any], _) -> Dict[str, Any]: +def format_search_describe_response(res: List[Any], _, __) -> Dict[str, Any]: """Format SEARCH.DESCRIBE response into index description.""" description: Dict[str, Any] = {} @@ -324,12 +325,23 @@ def format_search_describe_response(res: List[Any], _) -> Dict[str, Any]: return description -def format_search_count_response(res: Any, _) -> Dict[str, int]: +def format_search_count_response(res: Any, _, __) -> Dict[str, int]: """Format SEARCH.COUNT response.""" count = int(res) if not isinstance(res, int) else res return {"count": count} +def format_search_create_response(res: Any, command: List[str], client: Any): + """Format SEARCH.CREATE response by returning SearchIndexCommands instance.""" + # Extract index name from command (second parameter) + index_name = command[1] if len(command) > 1 else None + + if index_name is None: + return res + + return SearchIndexCommands(client, index_name) + + FORMATTERS: Dict[str, Callable] = { "COPY": to_bool, "EXPIRE": to_bool, @@ -405,19 +417,21 @@ def format_search_count_response(res: Any, _) -> Dict[str, int]: "XPENDING": format_xpending_response, "XGROUP DESTROY": to_bool, "XGROUP CREATECONSUMER": to_bool, + "SEARCH.CREATE": format_search_create_response, "SEARCH.QUERY": format_search_query_response, "SEARCH.DESCRIBE": format_search_describe_response, "SEARCH.COUNT": format_search_count_response, } -def cast_response(command: List[str], response: RESTResultT): +def cast_response(command: List[str], response: RESTResultT, client: Any): """ Given a command and its response, casts the response using the `FORMATTERS` map :param command: Used to determine the formatting to apply :param response: Response to format + :param client: Optional client instance for formatters that need to return client-dependent objects """ # get main command @@ -427,6 +441,6 @@ def cast_response(command: List[str], response: RESTResultT): # format response if main_command in FORMATTERS: - return FORMATTERS[main_command](response, command) + return FORMATTERS[main_command](response, command, client) return response From d1f737fd89d97629fc8ef98af47481f9cd76b7e4 Mon Sep 17 00:00:00 2001 From: CahidArda Date: Wed, 21 Jan 2026 16:49:47 +0300 Subject: [PATCH 05/12] fix: update return types of query and describe --- .../asyncio/search/test_search_index_async.py | 2 +- tests/commands/search/test_search_utils.py | 13 --- upstash_redis/commands.pyi | 10 +- upstash_redis/format.py | 96 ++----------------- upstash_redis/search_utils.py | 83 ---------------- 5 files changed, 12 insertions(+), 192 deletions(-) diff --git a/tests/commands/asyncio/search/test_search_index_async.py b/tests/commands/asyncio/search/test_search_index_async.py index c8e6146..3efa96e 100644 --- a/tests/commands/asyncio/search/test_search_index_async.py +++ b/tests/commands/asyncio/search/test_search_index_async.py @@ -283,7 +283,7 @@ async def test_count_with_numeric_filter(self, count_index): class TestAsyncDescribe: """Tests for async describe operations.""" - async def test_describe_index(self, async_redis_client, cleanup_indexes): + async def test_describe_index(self, async_redis_client: Redis, cleanup_indexes): """Test async describing an index.""" name = f"test-async-describe-{random_id()}" cleanup_indexes.append(name) diff --git a/tests/commands/search/test_search_utils.py b/tests/commands/search/test_search_utils.py index ec028ec..d1e1f09 100644 --- a/tests/commands/search/test_search_utils.py +++ b/tests/commands/search/test_search_utils.py @@ -8,7 +8,6 @@ build_query_command, deserialize_query_response, deserialize_describe_response, - parse_count_response, ) @@ -376,15 +375,3 @@ def test_deserialize_describe(self): assert "schema" in result assert "name" in result["schema"] assert result["schema"]["name"]["type"] == "TEXT" - - -class TestParseCountResponse: - """Tests for parse_count_response function.""" - - def test_parse_int_response(self): - """Test parsing integer count response.""" - assert parse_count_response(42) == 42 - - def test_parse_string_response(self): - """Test parsing string count response.""" - assert parse_count_response("42") == 42 diff --git a/upstash_redis/commands.pyi b/upstash_redis/commands.pyi index 3f06f74..cd91558 100644 --- a/upstash_redis/commands.pyi +++ b/upstash_redis/commands.pyi @@ -3,7 +3,7 @@ from typing import Any, Dict, List, Literal, Mapping, Optional, Tuple, Union from upstash_redis.typing import FloatMinMaxT, ValueT, JSONValueT from upstash_redis.utils import GeoSearchResult -from upstash_redis.search_types import NestedIndexSchema, FlatIndexSchema +from upstash_redis.search_types import NestedIndexSchema, FlatIndexSchema, IndexDescription, QueryResult class Commands: def bitcount( @@ -2052,7 +2052,7 @@ class SearchCommands: class SearchIndexCommands: def __init__(self, client: Commands, name: str, schema: Optional[Union[NestedIndexSchema, FlatIndexSchema]] = None): ... def wait_indexing(self) -> Any: ... - def describe(self) -> Dict[str, Any]: ... + def describe(self) -> IndexDescription: ... def query( self, *, @@ -2062,7 +2062,7 @@ class SearchIndexCommands: orderBy: Optional[Dict[str, str]] = None, select: Optional[Dict[str, bool]] = None, highlight: Optional[Dict[str, Any]] = None, - ) -> List[Dict[str, Any]]: ... + ) -> List[QueryResult]: ... def count(self, filter: Dict[str, Any]) -> Dict[str, int]: ... def drop(self) -> int: ... @@ -2088,7 +2088,7 @@ class AsyncSearchCommands: class AsyncSearchIndexCommands: def __init__(self, client: AsyncCommands, name: str, schema: Optional[Union[NestedIndexSchema, FlatIndexSchema]] = None): ... async def wait_indexing(self) -> Any: ... - async def describe(self) -> Dict[str, Any]: ... + async def describe(self) -> IndexDescription: ... async def query( self, *, @@ -2098,6 +2098,6 @@ class AsyncSearchIndexCommands: orderBy: Optional[Dict[str, str]] = None, select: Optional[Dict[str, bool]] = None, highlight: Optional[Dict[str, Any]] = None, - ) -> List[Dict[str, Any]]: ... + ) -> List[QueryResult]: ... async def count(self, filter: Dict[str, Any]) -> Dict[str, int]: ... async def drop(self) -> int: ... diff --git a/upstash_redis/format.py b/upstash_redis/format.py index 83f21a3..831976e 100644 --- a/upstash_redis/format.py +++ b/upstash_redis/format.py @@ -4,6 +4,10 @@ from upstash_redis.typing import RESTResultT from upstash_redis.utils import GeoSearchResult from upstash_redis.commands import SearchIndexCommands +from upstash_redis.search_utils import ( + deserialize_describe_response, + deserialize_query_response, +) def to_dict(res: List, _, __) -> Dict: @@ -229,100 +233,12 @@ def format_xpending_response(res, command, _): def format_search_query_response(res: List[Any], _, __) -> List[Dict[str, Any]]: """Format SEARCH.QUERY response into structured results.""" - results: List[Dict[str, Any]] = [] - - for item_raw in res: - if not isinstance(item_raw, (list, tuple)) or len(item_raw) < 2: - continue - - key = item_raw[0] - score = item_raw[1] - - result: Dict[str, Any] = {"key": key, "score": score} - - if len(item_raw) > 2: - raw_fields = item_raw[2] - - if isinstance(raw_fields, (list, tuple)) and len(raw_fields) > 0: - data: Dict[str, Any] = {} - - # Process field pairs - for field_raw in raw_fields: - if not isinstance(field_raw, (list, tuple)) or len(field_raw) < 2: - continue - - field_key = field_raw[0] - field_value = field_raw[1] - - # Handle nested paths - path_parts = field_key.split(".") - if len(path_parts) == 1: - data[field_key] = field_value - else: - # Build nested structure - current_obj = data - for i, part in enumerate(path_parts[:-1]): - if part not in current_obj: - current_obj[part] = {} - current_obj = current_obj[part] - current_obj[path_parts[-1]] = field_value - - # If $ key exists (full document), use its contents - if "$" in data: - data = data["$"] - - result["data"] = data - - results.append(result) - - return results + return deserialize_query_response(res) def format_search_describe_response(res: List[Any], _, __) -> Dict[str, Any]: """Format SEARCH.DESCRIBE response into index description.""" - description: Dict[str, Any] = {} - - i = 0 - while i < len(res): - descriptor = res[i] - value = res[i + 1] if i + 1 < len(res) else None - - if descriptor == "name": - description["name"] = value - elif descriptor == "type": - if value: - description["dataType"] = value.lower() - elif descriptor == "prefixes": - description["prefixes"] = value if isinstance(value, list) else [] - elif descriptor == "language": - description["language"] = value - elif descriptor == "schema": - schema: Dict[str, Dict[str, Any]] = {} - if isinstance(value, list): - for field_desc in value: - if not isinstance(field_desc, list) or len(field_desc) < 2: - continue - - field_name = field_desc[0] - field_info: Dict[str, Any] = {"type": field_desc[1]} - - # Parse field options - for j in range(2, len(field_desc)): - option = field_desc[j] - if option == "NOSTEM": - field_info["noStem"] = True - elif option == "NOTOKENIZE": - field_info["noTokenize"] = True - elif option == "FAST": - field_info["fast"] = True - - schema[field_name] = field_info - - description["schema"] = schema - - i += 2 - - return description + return deserialize_describe_response(res) def format_search_count_response(res: Any, _, __) -> Dict[str, int]: diff --git a/upstash_redis/search_utils.py b/upstash_redis/search_utils.py index 4c32b28..2320c55 100644 --- a/upstash_redis/search_utils.py +++ b/upstash_redis/search_utils.py @@ -155,75 +155,6 @@ def build_create_index_command(params: CreateIndexParams) -> List[str]: return command -def _build_filter_query(filter_dict: Dict[str, Any]) -> List[str]: - """ - Build filter query from filter dictionary. - - Args: - filter_dict: Filter specification - - Returns: - List of filter query parts - """ - filters: List[str] = [] - - for field_name, field_filter in filter_dict.items(): - if not isinstance(field_filter, dict): - continue - - # Text filters - if "eq" in field_filter: - value = field_filter["eq"] - if isinstance(value, bool): - filters.append(f"@{field_name}:[{str(value).lower()}]") - else: - filters.append(f"@{field_name}:({value})") - - elif "fuzzy" in field_filter: - fuzzy = field_filter["fuzzy"] - if isinstance(fuzzy, str): - filters.append(f"@{field_name}:%{fuzzy}%") - elif isinstance(fuzzy, dict): - value = fuzzy.get("value", "") - distance = fuzzy.get("distance", 1) - filters.append(f"@{field_name}:%" + "%" * distance + value + "%" * distance) - - elif "phrase" in field_filter: - phrase = field_filter["phrase"] - if isinstance(phrase, str): - filters.append(f'@{field_name}:"{phrase}"') - elif isinstance(phrase, dict): - value = phrase.get("value", "") - prefix = phrase.get("prefix", False) - if prefix: - filters.append(f'@{field_name}:"{value}"*') - else: - filters.append(f'@{field_name}:"{value}"') - - elif "regex" in field_filter: - regex = field_filter["regex"] - filters.append(f"@{field_name}:/{regex}/") - - # Numeric filters - elif "gt" in field_filter or "gte" in field_filter or "lt" in field_filter or "lte" in field_filter: - min_val = "-inf" - max_val = "+inf" - - if "gt" in field_filter: - min_val = f"({field_filter['gt']}" - elif "gte" in field_filter: - min_val = str(field_filter["gte"]) - - if "lt" in field_filter: - max_val = f"({field_filter['lt']}" - elif "lte" in field_filter: - max_val = str(field_filter["lte"]) - - filters.append(f"@{field_name}:[{min_val} {max_val}]") - - return filters - - def build_query_command( command_name: str, index_name: str, options: QueryOptions = None ) -> List[str]: @@ -420,17 +351,3 @@ def deserialize_describe_response(raw_response: List[Any]) -> IndexDescription: return description - -def parse_count_response(raw_response: Any) -> int: - """ - Parse count response. - - Args: - raw_response: Raw response from SEARCH.COUNT - - Returns: - Count as integer - """ - if isinstance(raw_response, int): - return raw_response - return int(raw_response) From 95c24636b5c1f75cf873ce0d54aba59b93df4a4b Mon Sep 17 00:00:00 2001 From: CahidArda Date: Wed, 21 Jan 2026 17:14:50 +0300 Subject: [PATCH 06/12] feat: check query results in more detail in tests --- .../asyncio/search/test_search_index_async.py | 183 +++++++++++++++--- tests/commands/search/test_search_index.py | 169 +++++++++++++--- upstash_redis/commands.pyi | 4 +- upstash_redis/search_utils.py | 12 +- 4 files changed, 314 insertions(+), 54 deletions(-) diff --git a/tests/commands/asyncio/search/test_search_index_async.py b/tests/commands/asyncio/search/test_search_index_async.py index 3efa96e..ac16495 100644 --- a/tests/commands/asyncio/search/test_search_index_async.py +++ b/tests/commands/asyncio/search/test_search_index_async.py @@ -3,12 +3,21 @@ import json import random import string -from typing import Any, Dict, List +from typing import Any, Dict, List, TypedDict import pytest import pytest_asyncio from upstash_redis.asyncio import Redis +from upstash_redis.commands import AsyncSearchIndexCommands + + +class QueryIndexFixture(TypedDict): + """Type definition for query_index fixture.""" + index: AsyncSearchIndexCommands + redis: Redis + keys: List[str] + name: str def random_id() -> str: @@ -115,69 +124,106 @@ async def query_index(self, async_redis_client: Redis, cleanup_indexes): if keys: await redis.delete(*keys) - async def test_query_text_eq(self, query_index): + async def test_query_text_eq(self, query_index: QueryIndexFixture): """Test async querying with $eq on text field.""" index = query_index["index"] result = await index.query(filter={"name": {"$eq": "Laptop"}}) - assert len(result) > 0 - assert all("key" in r and "score" in r for r in result) + assert len(result) == 2 + for r in result: + assert "key" in r + # Verify key has the expected prefix + assert query_index["name"] in r["key"] + # Verify score is numeric + assert "score" in r + assert isinstance(float(r["score"]), float) - async def test_query_text_fuzzy(self, query_index): + async def test_query_text_fuzzy(self, query_index: QueryIndexFixture): """Test async querying with fuzzy search.""" index = query_index["index"] result = await index.query(filter={"name": {"$fuzzy": "laptopp"}}) - assert len(result) > 0 + # Fuzzy search should find Laptop variants despite typo + assert len(result) >= 2 + assert all("key" in r and "score" in r for r in result) - async def test_query_text_phrase(self, query_index): + async def test_query_text_phrase(self, query_index: QueryIndexFixture): """Test async querying with phrase matching.""" index = query_index["index"] result = await index.query(filter={"description": {"$phrase": "wireless mouse"}}) - assert len(result) > 0 + # Should find the Wireless Mouse item + assert len(result) == 1 + assert "key" in result[0] + assert "score" in result[0] + assert isinstance(float(result[0]["score"]), float) + # Check data field exists and has expected fields + if "data" in result[0]: + assert "name" in result[0]["data"] or "description" in result[0]["data"] - async def test_query_numeric_gt(self, query_index): + async def test_query_numeric_gt(self, query_index: QueryIndexFixture): """Test async querying with $gt on numeric field.""" index = query_index["index"] result = await index.query(filter={"price": {"$gt": 500}}) - assert len(result) > 0 - # Verify data structure + # Should find Laptop Pro (1299.99) and Laptop Basic (599.99) + assert len(result) == 2 for item in result: assert "key" in item assert "score" in item - - async def test_query_numeric_gte(self, query_index): + assert isinstance(float(item["score"]), float) + # Check data field exists with price field + if "data" in item: + assert "price" in item["data"] + # Verify price is greater than 500 + price_value = item["data"]["price"] + assert isinstance(price_value, (int, float, str)) + assert float(price_value) > 500 + + async def test_query_numeric_gte(self, query_index: QueryIndexFixture): """Test async querying with $gte on numeric field.""" index = query_index["index"] result = await index.query(filter={"stock": {"$gte": 100}}) - assert len(result) > 0 + # Should find items with stock >= 100 (Laptop Basic, Mouse, Cable) + assert len(result) == 3 + assert all("key" in r and "score" in r for r in result) - async def test_query_numeric_lt(self, query_index): + async def test_query_numeric_lt(self, query_index: QueryIndexFixture): """Test async querying with $lt on numeric field.""" index = query_index["index"] result = await index.query(filter={"price": {"$lt": 50}}) - assert len(result) > 0 + # Should find Wireless Mouse (29.99) and USB Cable (9.99) + assert len(result) == 2 + assert all("key" in r and "score" in r for r in result) - async def test_query_boolean_eq(self, query_index): + async def test_query_boolean_eq(self, query_index: QueryIndexFixture): """Test async querying with $eq on boolean field.""" index = query_index["index"] result = await index.query(filter={"active": {"$eq": False}}) - assert len(result) > 0 - - async def test_query_with_limit(self, query_index): + # Should find USB Cable (active: False) + assert len(result) == 1 + assert "key" in result[0] + assert "score" in result[0] + assert isinstance(float(result[0]["score"]), float) + # Check data field has active field + if "data" in result[0]: + assert "active" in result[0]["data"] + # Verify active is False + assert result[0]["data"]["active"] in (False, "false", "False", 0, "0") + + async def test_query_with_limit(self, query_index: QueryIndexFixture): """Test async querying with limit option.""" index = query_index["index"] result = await index.query(filter={"active": {"$eq": True}}, limit=2) - assert len(result) > 0 - assert len(result) <= 2 + # Limit should restrict results to 2 even though 3 items match + assert len(result) == 2 + assert all("key" in r and "score" in r for r in result) - async def test_query_with_pagination(self, query_index): + async def test_query_with_pagination(self, query_index: QueryIndexFixture): """Test async querying with offset for pagination.""" index = query_index["index"] first_page = await index.query( @@ -187,9 +233,16 @@ async def test_query_with_pagination(self, query_index): filter={"active": {"$eq": True}}, limit=2, offset=2 ) - assert len(first_page) > 0 + # Should get 2 items on first page + assert len(first_page) == 2 + # Should get 1 item on second page (3 total active items) + assert len(second_page) == 1 + # Pages should have different items + first_keys = {r["key"] for r in first_page} + second_keys = {r["key"] for r in second_page} + assert first_keys.isdisjoint(second_keys) - async def test_query_with_sort_asc(self, query_index): + async def test_query_with_sort_asc(self, query_index: QueryIndexFixture): """Test async querying with sortBy ascending.""" index = query_index["index"] result = await index.query( @@ -198,9 +251,11 @@ async def test_query_with_sort_asc(self, query_index): limit=3, ) - assert len(result) > 0 + # Should get 3 active items sorted by price ascending + assert len(result) == 3 + assert all("key" in r and "score" in r for r in result) - async def test_query_with_sort_desc(self, query_index): + async def test_query_with_sort_desc(self, query_index: QueryIndexFixture): """Test async querying with sortBy descending.""" index = query_index["index"] result = await index.query( @@ -209,7 +264,79 @@ async def test_query_with_sort_desc(self, query_index): limit=3, ) - assert len(result) > 0 + # Should get 3 active items sorted by price descending + assert len(result) == 3 + assert all("key" in r and "score" in r for r in result) + + async def test_query_with_all_fields(self, query_index: QueryIndexFixture): + """Test querying returns all fields when select is not specified.""" + index = query_index["index"] + result = await index.query(filter={"name": {"$eq": "Wireless Mouse"}}, limit=1) + + # Should get complete result with all fields + assert len(result) == 1 + item = result[0] + + # Verify structure + assert "key" in item + assert "score" in item + assert isinstance(float(item["score"]), float) + assert "data" in item + + # Verify all fields are present in data (values are strings for string dataType) + data = item["data"] + assert data["name"] == "Wireless Mouse" + assert data["description"] == "Ergonomic wireless mouse" + assert float(data["price"]) == 29.99 + assert int(data["stock"]) == 200 + assert data["active"] in (True, "true", "True") + + async def test_query_with_no_content(self, query_index: QueryIndexFixture): + """Test querying with select={{}} returns only keys and scores.""" + index = query_index["index"] + result = await index.query( + filter={"name": {"$eq": "Wireless Mouse"}}, + select={}, + limit=1 + ) + + # Should get only key and score, no data + assert len(result) == 1 + item = result[0] + + assert "key" in item + assert "score" in item + assert isinstance(float(item["score"]), float) + # With noContent, data field should not be present + assert "data" not in item + + async def test_query_with_specific_field(self, query_index: QueryIndexFixture): + """Test querying with select={{"price": True}} returns only specified field.""" + index = query_index["index"] + result = await index.query( + filter={"name": {"$eq": "Wireless Mouse"}}, + select={"price": True}, + limit=1 + ) + + # Should get key, score, and only the price field + assert len(result) == 1 + item = result[0] + + assert "key" in item + assert "score" in item + assert isinstance(float(item["score"]), float) + assert "data" in item + + # Verify only price field is present (value is string for string dataType) + data = item["data"] + assert "price" in data + assert float(data["price"]) == 29.99 + # Other fields should not be present + assert "name" not in data + assert "description" not in data + assert "stock" not in data + assert "active" not in data @pytest.mark.asyncio diff --git a/tests/commands/search/test_search_index.py b/tests/commands/search/test_search_index.py index 9abe34e..1caa854 100644 --- a/tests/commands/search/test_search_index.py +++ b/tests/commands/search/test_search_index.py @@ -259,71 +259,127 @@ def test_query_text_eq(self, string_index: StringIndexFixture): index = string_index["index"] result = index.query(filter={"name": {"$eq": "Laptop"}}) - assert len(result) > 0 + # Should find 2 laptops: Laptop Pro and Laptop Basic + assert len(result) == 2 + for r in result: + assert "key" in r + # Verify key has the expected prefix + assert string_index["name"] in r["key"] + # Verify score is numeric + assert "score" in r + assert isinstance(float(r["score"]), float) def test_query_text_fuzzy(self, string_index: StringIndexFixture): """Test querying with fuzzy search for typo tolerance.""" index = string_index["index"] result = index.query(filter={"name": {"$fuzzy": "laptopp"}}) - assert len(result) > 0 + # Fuzzy search should find Laptop variants despite typo + assert len(result) >= 2 + assert all("key" in r and "score" in r for r in result) def test_query_text_phrase(self, string_index: StringIndexFixture): """Test querying with phrase matching.""" index = string_index["index"] result = index.query(filter={"description": {"$phrase": "wireless mouse"}}) - assert len(result) > 0 + # Should find exactly the Wireless Mouse item + assert len(result) == 1 + assert "key" in result[0] + assert "score" in result[0] + assert isinstance(float(result[0]["score"]), float) + # Check data field exists + if "data" in result[0]: + # Check description field in data + assert "description" in result[0]["data"] or "name" in result[0]["data"] def test_query_text_regex(self, string_index: StringIndexFixture): """Test querying with regex pattern.""" index = string_index["index"] result = index.query(filter={"name": {"$regex": "Laptop.*"}}) - assert len(result) > 0 + # Should find both Laptop Pro and Laptop Basic + assert len(result) == 2 + assert all("key" in r and "score" in r for r in result) def test_query_numeric_gt(self, string_index: StringIndexFixture): """Test querying with $gt on numeric field.""" index = string_index["index"] result = index.query(filter={"price": {"$gt": 500}}) - assert len(result) > 0 + # Should find Laptop Pro (1299.99) and Laptop Basic (599.99) + assert len(result) == 2 + for r in result: + assert "key" in r + assert "score" in r + assert isinstance(float(r["score"]), float) + # Check data field with price + if "data" in r: + assert "price" in r["data"] + # Verify price is greater than 500 + price_value = r["data"]["price"] + assert isinstance(price_value, (int, float, str)) + assert float(price_value) > 500 def test_query_numeric_gte(self, string_index: StringIndexFixture): """Test querying with $gte on numeric field.""" index = string_index["index"] result = index.query(filter={"stock": {"$gte": 100}}) - assert len(result) > 0 + # Should find items with stock >= 100 (Laptop Basic: 100, Mouse: 200, Cable: 500, Case: 300) + assert len(result) == 4 + assert all("key" in r and "score" in r for r in result) def test_query_numeric_lt(self, string_index: StringIndexFixture): """Test querying with $lt on numeric field.""" index = string_index["index"] result = index.query(filter={"price": {"$lt": 50}}) - assert len(result) > 0 + # Should find Wireless Mouse (29.99), USB Cable (9.99), Phone Case (19.99) + assert len(result) == 3 + assert all("key" in r and "score" in r for r in result) def test_query_numeric_lte(self, string_index: StringIndexFixture): """Test querying with $lte on numeric field.""" index = string_index["index"] result = index.query(filter={"stock": {"$lte": 50}}) - assert len(result) > 0 + # Should find only Laptop Pro (stock: 50) + assert len(result) == 1 + assert "key" in result[0] + assert "score" in result[0] + assert isinstance(float(result[0]["score"]), float) + # Check data field with stock + if "data" in result[0]: + assert "stock" in result[0]["data"] + stock_value = result[0]["data"]["stock"] + assert isinstance(stock_value, (int, float, str)) + assert float(stock_value) <= 50 def test_query_boolean_eq(self, string_index: StringIndexFixture): """Test querying with $eq on boolean field.""" index = string_index["index"] result = index.query(filter={"active": {"$eq": False}}) - assert len(result) > 0 + # Should find only Phone Case (active: False) + assert len(result) == 1 + assert "key" in result[0] + assert "score" in result[0] + assert isinstance(float(result[0]["score"]), float) + # Check data field with active + if "data" in result[0]: + assert "active" in result[0]["data"] + # Verify active is False + assert result[0]["data"]["active"] in (False, "false", "False", 0, "0") def test_query_with_limit(self, string_index: StringIndexFixture): """Test querying with limit option.""" index = string_index["index"] result = index.query(filter={"category": {"$eq": "electronics"}}, limit=2) - assert len(result) > 0 - assert len(result) <= 2 + # Should limit to exactly 2 results even though 3 match + assert len(result) == 2 + assert all("key" in r and "score" in r for r in result) def test_query_with_pagination(self, string_index: StringIndexFixture): """Test querying with offset for pagination.""" @@ -336,8 +392,14 @@ def test_query_with_pagination(self, string_index: StringIndexFixture): filter={"category": {"$eq": "electronics"}}, limit=2, offset=2 ) - assert len(first_page) > 0 - # second_page might be empty if there are only 2-3 results + # Should get 2 items on first page + assert len(first_page) == 2 + # Should get 1 item on second page (3 total electronics) + assert len(second_page) == 1 + # Pages should have different items + first_keys = {r["key"] for r in first_page} + second_keys = {r["key"] for r in second_page} + assert first_keys.isdisjoint(second_keys) def test_query_with_sort_asc(self, string_index: StringIndexFixture): """Test querying with sortBy ascending.""" @@ -348,7 +410,9 @@ def test_query_with_sort_asc(self, string_index: StringIndexFixture): limit=3, ) - assert len(result) > 0 + # Should get all 3 electronics items sorted by price + assert len(result) == 3 + assert all("key" in r and "score" in r for r in result) def test_query_with_sort_desc(self, string_index: StringIndexFixture): """Test querying with sortBy descending.""" @@ -359,25 +423,82 @@ def test_query_with_sort_desc(self, string_index: StringIndexFixture): limit=3, ) - assert len(result) > 0 + # Should get all 3 electronics items sorted by price descending + assert len(result) == 3 + assert all("key" in r and "score" in r for r in result) + + def test_query_with_all_fields(self, string_index: StringIndexFixture): + """Test querying returns all fields when select is not specified.""" + index = string_index["index"] + result = index.query(filter={"name": {"$eq": "Wireless Mouse"}}, limit=1) + + # Should get complete result with all fields + assert len(result) == 1 + item = result[0] + + # Verify structure + assert "key" in item + assert "score" in item + assert isinstance(float(item["score"]), float) + assert "data" in item + + # Verify all fields are present in data (values are strings for string dataType) + data = item["data"] + assert data["name"] == "Wireless Mouse" + assert data["description"] == "Ergonomic wireless mouse" + assert data["category"] == "electronics" + assert float(data["price"]) == 29.99 + assert int(data["stock"]) == 200 + assert data["active"] in (True, "true", "True") def test_query_no_content(self, string_index: StringIndexFixture): """Test querying with noContent (keys only).""" index = string_index["index"] - result = index.query(filter={"category": {"$eq": "electronics"}}, select={}) + result = index.query( + filter={"name": {"$eq": "Wireless Mouse"}}, + select={}, + limit=1 + ) - assert len(result) > 0 + # Should get only key and score, no data field + assert len(result) == 1 + item = result[0] + + assert "key" in item + assert "score" in item + assert isinstance(float(item["score"]), float) + # With noContent, data field should not be present + assert "data" not in item def test_query_with_return_fields(self, string_index: StringIndexFixture): """Test querying with specific return fields.""" index = string_index["index"] result = index.query( - filter={"category": {"$eq": "electronics"}}, - select={"category": True}, - highlight={"fields": []}, + filter={"name": {"$eq": "Wireless Mouse"}}, + select={"price": True, "stock": True}, + limit=1 ) - assert len(result) > 0 + # Should get only specified fields + assert len(result) == 1 + item = result[0] + + assert "key" in item + assert "score" in item + assert isinstance(float(item["score"]), float) + assert "data" in item + + # Verify only selected fields are present (values are strings for string dataType) + data = item["data"] + assert "price" in data + assert float(data["price"]) == 29.99 + assert "stock" in data + assert int(data["stock"]) == 200 + # Other fields should not be present + assert "name" not in data + assert "description" not in data + assert "category" not in data + assert "active" not in data class TestCount: @@ -422,7 +543,8 @@ def test_count_matching_documents(self, count_index: CountIndexFixture): result = index.count({"type": {"$eq": "A"}}) assert "count" in result - assert result["count"] > 0 + # Should count exactly 5 documents of type A (indices 0-4) + assert result["count"] == 5 def test_count_with_numeric_filter(self, count_index: CountIndexFixture): """Test counting with numeric filter.""" @@ -430,7 +552,8 @@ def test_count_with_numeric_filter(self, count_index: CountIndexFixture): result = index.count({"value": {"$eq": 5}}) assert "count" in result - assert result["count"] > 0 + # Should count exactly 1 document with value 5 + assert result["count"] == 1 class TestDescribe: diff --git a/upstash_redis/commands.pyi b/upstash_redis/commands.pyi index cd91558..aaf4c16 100644 --- a/upstash_redis/commands.pyi +++ b/upstash_redis/commands.pyi @@ -2063,7 +2063,7 @@ class SearchIndexCommands: select: Optional[Dict[str, bool]] = None, highlight: Optional[Dict[str, Any]] = None, ) -> List[QueryResult]: ... - def count(self, filter: Dict[str, Any]) -> Dict[str, int]: ... + def count(self, filter: Dict[str, Any]) -> Dict[Literal["count"], int]: ... def drop(self) -> int: ... class AsyncSearchCommands: @@ -2099,5 +2099,5 @@ class AsyncSearchIndexCommands: select: Optional[Dict[str, bool]] = None, highlight: Optional[Dict[str, Any]] = None, ) -> List[QueryResult]: ... - async def count(self, filter: Dict[str, Any]) -> Dict[str, int]: ... + async def count(self, filter: Dict[str, Any]) -> Dict[Literal["count"], int]: ... async def drop(self) -> int: ... diff --git a/upstash_redis/search_utils.py b/upstash_redis/search_utils.py index 2320c55..61774f2 100644 --- a/upstash_redis/search_utils.py +++ b/upstash_redis/search_utils.py @@ -288,7 +288,17 @@ def deserialize_query_response(raw_response: List[Any]) -> List[QueryResult]: # If $ key exists (full document), use its contents if "$" in data: - data = data["$"] + dollar_value = data["$"] + # Parse JSON string if needed + if isinstance(dollar_value, str): + import json + try: + data = json.loads(dollar_value) + except (json.JSONDecodeError, ValueError): + # If parsing fails, use the string as-is + data = dollar_value + else: + data = dollar_value result["data"] = data From 406154824ebc1912d0d6b5daf0989c6beaf48894 Mon Sep 17 00:00:00 2001 From: CahidArda Date: Wed, 21 Jan 2026 17:27:57 +0300 Subject: [PATCH 07/12] feat: add example --- examples/search.py | 334 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 334 insertions(+) create mode 100644 examples/search.py diff --git a/examples/search.py b/examples/search.py new file mode 100644 index 0000000..620c168 --- /dev/null +++ b/examples/search.py @@ -0,0 +1,334 @@ +""" +Example demonstrating Redis Search functionality with the Upstash Python SDK. + +This example shows how to: +- Create search indexes for different data types (string, JSON, hash) +- Index data with various field types (TEXT, numeric, boolean) +- Query with filters, sorting, and field selection +- Work with nested schemas +""" + +import json +import dotenv +from upstash_redis import Redis + +dotenv.load_dotenv() + +# Initialize Redis client +redis = Redis.from_env() + +# ============================================================================= +# Example 1: Simple String Index with Product Catalog +# ============================================================================= + +print("=" * 60) +print("Example 1: String Index - Product Catalog") +print("=" * 60) + +# Create an index for products stored as JSON strings +products_index = redis.search.create_index( + name="products", + dataType="string", # Data is stored as strings + prefix="product:", # Index keys starting with "product:" + schema={ + "name": "TEXT", # Full-text searchable + "category": {"type": "TEXT", "noTokenize": True}, # Exact match only + "price": {"type": "F64", "fast": True}, # Float, optimized for sorting + "stock": {"type": "U64", "fast": True}, # Unsigned integer + "active": "BOOL", # Boolean field + } +) + +# Add product data +products = [ + { + "name": "Laptop Pro", + "category": "electronics", + "price": 1299.99, + "stock": 50, + "active": True, + }, + { + "name": "Laptop Basic", + "category": "electronics", + "price": 599.99, + "stock": 100, + "active": True, + }, + { + "name": "Wireless Mouse", + "category": "electronics", + "price": 29.99, + "stock": 200, + "active": True, + }, + { + "name": "USB Cable", + "category": "accessories", + "price": 9.99, + "stock": 500, + "active": False, + }, +] + +for i, product in enumerate(products): + redis.set(f"product:{i}", json.dumps(product)) + +# Wait for indexing to complete +products_index.wait_indexing() + +# Query 1: Search for laptops +print("\n1. Search for 'Laptop' in name:") +results = products_index.query(filter={"name": {"$eq": "Laptop"}}) +for result in results: + print(f" Key: {result['key']}, Score: {result['score']}") + print(f" Data: {result['data']}") + +# Query 2: Find products over $500 +print("\n2. Find products with price > 500:") +results = products_index.query(filter={"price": {"$gt": 500}}) +for result in results: + data = result["data"] + print(f" {data['name']}: ${data['price']}") + +# Query 3: Get inactive products +print("\n3. Find inactive products:") +results = products_index.query(filter={"active": {"$eq": False}}) +for result in results: + data = result["data"] + print(f" {data['name']} (stock: {data['stock']})") + +# Query 4: Search with sorting and field selection +print("\n4. Electronics sorted by price (descending), show only name and price:") +results = products_index.query( + filter={"category": {"$eq": "electronics"}}, + orderBy={"price": "DESC"}, + select={"name": True, "price": True}, +) +for result in results: + data = result["data"] + print(f" {data['name']}: ${data['price']}") + +# Query 5: Pagination +print("\n5. Pagination example (limit=2, offset=0 and offset=2):") +page1 = products_index.query( + filter={"category": {"$eq": "electronics"}}, + limit=2, + offset=0, +) +print(f" Page 1: {len(page1)} results") +for result in page1: + print(f" - {result['data']['name']}") + +page2 = products_index.query( + filter={"category": {"$eq": "electronics"}}, + limit=2, + offset=2, +) +print(f" Page 2: {len(page2)} results") +for result in page2: + print(f" - {result['data']['name']}") + +# Query 6: Count documents +print("\n6. Count electronics products:") +count_result = products_index.count(filter={"category": {"$eq": "electronics"}}) +print(f" Total electronics: {count_result['count']}") + +# Query 7: Get only keys and scores (no content) +print("\n7. Get keys only (select={}):") +results = products_index.query( + filter={"category": {"$eq": "electronics"}}, + select={}, + limit=2, +) +for result in results: + print(f" Key: {result['key']}, Score: {result['score']}") + print(f" Has data field: {'data' in result}") + +# Describe the index +print("\n8. Index description:") +description = products_index.describe() +print(f" Name: {description['name']}") +print(f" Type: {description['dataType']}") +print(f" Prefixes: {description['prefixes']}") +print(f" Schema fields: {list(description['schema'].keys())}") + +# ============================================================================= +# Example 2: JSON Index with Nested Schema +# ============================================================================= + +print("\n" + "=" * 60) +print("Example 2: JSON Index - Blog Posts with Nested Data") +print("=" * 60) + +# Create index with nested schema +posts_index = redis.search.create_index( + name="posts", + dataType="json", # Data is stored as JSON + prefix="post:", + schema={ + "title": "TEXT", + "author": { + "name": "TEXT", + "email": "TEXT", + }, + "stats": { + "views": {"type": "U64", "fast": True}, + "likes": {"type": "U64", "fast": True}, + }, + "published": "BOOL", + } +) + +# Add blog posts +posts = [ + { + "title": "Getting Started with Redis", + "author": {"name": "John Doe", "email": "john@example.com"}, + "stats": {"views": 1500, "likes": 75}, + "published": True, + }, + { + "title": "Advanced Redis Patterns", + "author": {"name": "Jane Smith", "email": "jane@example.com"}, + "stats": {"views": 800, "likes": 40}, + "published": True, + }, + { + "title": "Redis Search Tutorial", + "author": {"name": "John Doe", "email": "john@example.com"}, + "stats": {"views": 2000, "likes": 120}, + "published": True, + }, +] + +for i, post in enumerate(posts): + redis.json.set(f"post:{i}", "$", post) + +posts_index.wait_indexing() + +# Query nested fields +print("\n1. Posts by author 'John':") +results = posts_index.query(filter={"author.name": {"$eq": "John"}}) +for result in results: + data = result["data"] + print(f" Title: {data['title']}") + print(f" Author: {data['author']['name']}") + print(f" Views: {data['stats']['views']}") + +# Query with nested numeric field +print("\n2. Posts with more than 1000 views:") +results = posts_index.query( + filter={"stats.views": {"$gt": 1000}}, + orderBy={"stats.views": "DESC"}, +) +for result in results: + data = result["data"] + print(f" {data['title']}: {data['stats']['views']} views") + +# Select nested fields +print("\n3. Get only author email and views count:") +results = posts_index.query( + filter={"published": {"$eq": True}}, + select={"author.email": True, "stats.views": True}, + limit=2, +) +for result in results: + data = result["data"] + print(f" Email: {data['author']['email']}, Views: {data['stats']['views']}") + +# ============================================================================= +# Example 3: Hash Index +# ============================================================================= + +print("\n" + "=" * 60) +print("Example 3: Hash Index - User Scores") +print("=" * 60) + +# Create hash index +scores_index = redis.search.create_index( + name="scores", + dataType="hash", # Data is stored as Redis hash + prefix="user:", + schema={ + "username": "TEXT", + "score": {"type": "U64", "fast": True}, + "level": {"type": "U64", "fast": True}, + } +) + +# Add user data using HSET +users = [ + {"username": "alice", "score": "9500", "level": "10"}, + {"username": "bob", "score": "8700", "level": "9"}, + {"username": "charlie", "score": "9200", "level": "10"}, +] + +for i, user in enumerate(users): + redis.hset(f"user:{i}", values=user) + +scores_index.wait_indexing() + +# Query hash data +print("\n1. Top level 10 players:") +results = scores_index.query( + filter={"level": {"$eq": 10}}, + orderBy={"score": "DESC"}, +) +for result in results: + data = result["data"] + print(f" {data['username']}: {data['score']} points") + +# ============================================================================= +# Example 4: Advanced Queries +# ============================================================================= + +print("\n" + "=" * 60) +print("Example 4: Advanced Query Features") +print("=" * 60) + +# Fuzzy search (typo tolerance) +print("\n1. Fuzzy search for 'laptopp' (with typo):") +results = products_index.query(filter={"name": {"$fuzzy": "laptopp"}}) +print(f" Found {len(results)} results despite typo") +for result in results: + print(f" - {result['data']['name']}") + +# Phrase search +print("\n2. Phrase search for 'Wireless Mouse':") +results = products_index.query(filter={"name": {"$phrase": "Wireless Mouse"}}) +for result in results: + print(f" Found: {result['data']['name']}") + +# Range query +print("\n3. Products priced between $10 and $100:") +results = products_index.query( + filter={"price": {"$gte": 10, "$lte": 100}}, + orderBy={"price": "ASC"}, +) +for result in results: + data = result["data"] + print(f" {data['name']}: ${data['price']}") + +# ============================================================================= +# Cleanup +# ============================================================================= + +print("\n" + "=" * 60) +print("Cleanup") +print("=" * 60) + +# Drop indexes +products_index.drop() +posts_index.drop() +scores_index.drop() + +# Delete data +for i in range(len(products)): + redis.delete(f"product:{i}") +for i in range(len(posts)): + redis.delete(f"post:{i}") +for i in range(len(users)): + redis.delete(f"user:{i}") + +print("All indexes and data cleaned up!") From 6f1bce49f8663c9916f1ea2e0a3da596bb5cde35 Mon Sep 17 00:00:00 2001 From: Metin Dumandag <29387993+mdumandag@users.noreply.github.com> Date: Fri, 23 Jan 2026 13:00:21 +0300 Subject: [PATCH 08/12] cleanup --- examples/search.py | 84 ++-- pyproject.toml | 5 +- .../asyncio/search/test_search_index_async.py | 192 ++++----- tests/commands/search/test_search_advanced.py | 81 ++-- tests/commands/search/test_search_index.py | 291 +++++++------- tests/commands/search/test_search_utils.py | 377 ------------------ tests/test_formatters.py | 126 +++--- upstash_redis/asyncio/client.py | 4 +- upstash_redis/client.py | 4 +- upstash_redis/commands.py | 220 ++++++---- upstash_redis/commands.pyi | 66 +-- upstash_redis/format.py | 20 +- upstash_redis/search.py | 220 ++++++++++ upstash_redis/search_types.py | 115 ------ upstash_redis/search_utils.py | 363 ----------------- 15 files changed, 775 insertions(+), 1393 deletions(-) delete mode 100644 tests/commands/search/test_search_utils.py create mode 100644 upstash_redis/search.py delete mode 100644 upstash_redis/search_types.py delete mode 100644 upstash_redis/search_utils.py diff --git a/examples/search.py b/examples/search.py index 620c168..bf8614b 100644 --- a/examples/search.py +++ b/examples/search.py @@ -9,7 +9,9 @@ """ import json + import dotenv + from upstash_redis import Redis dotenv.load_dotenv() @@ -28,15 +30,15 @@ # Create an index for products stored as JSON strings products_index = redis.search.create_index( name="products", - dataType="string", # Data is stored as strings - prefix="product:", # Index keys starting with "product:" + data_type="string", # Data is stored as strings + prefixes="product:", # Index keys starting with "product:" schema={ "name": "TEXT", # Full-text searchable - "category": {"type": "TEXT", "noTokenize": True}, # Exact match only + "category": {"type": "TEXT", "no_tokenize": True}, # Exact match only "price": {"type": "F64", "fast": True}, # Float, optimized for sorting "stock": {"type": "U64", "fast": True}, # Unsigned integer "active": "BOOL", # Boolean field - } + }, ) # Add product data @@ -81,32 +83,32 @@ print("\n1. Search for 'Laptop' in name:") results = products_index.query(filter={"name": {"$eq": "Laptop"}}) for result in results: - print(f" Key: {result['key']}, Score: {result['score']}") - print(f" Data: {result['data']}") + print(f" Key: {result.key}, Score: {result.score}") + print(f" Data: {result.data}") # Query 2: Find products over $500 print("\n2. Find products with price > 500:") results = products_index.query(filter={"price": {"$gt": 500}}) for result in results: - data = result["data"] + data = result.data print(f" {data['name']}: ${data['price']}") # Query 3: Get inactive products print("\n3. Find inactive products:") results = products_index.query(filter={"active": {"$eq": False}}) for result in results: - data = result["data"] + data = result.data print(f" {data['name']} (stock: {data['stock']})") # Query 4: Search with sorting and field selection print("\n4. Electronics sorted by price (descending), show only name and price:") results = products_index.query( filter={"category": {"$eq": "electronics"}}, - orderBy={"price": "DESC"}, + order_by={"price": "DESC"}, select={"name": True, "price": True}, ) for result in results: - data = result["data"] + data = result.data print(f" {data['name']}: ${data['price']}") # Query 5: Pagination @@ -118,7 +120,7 @@ ) print(f" Page 1: {len(page1)} results") for result in page1: - print(f" - {result['data']['name']}") + print(f" - {result.data['name']}") page2 = products_index.query( filter={"category": {"$eq": "electronics"}}, @@ -127,12 +129,12 @@ ) print(f" Page 2: {len(page2)} results") for result in page2: - print(f" - {result['data']['name']}") + print(f" - {result.data['name']}") # Query 6: Count documents print("\n6. Count electronics products:") count_result = products_index.count(filter={"category": {"$eq": "electronics"}}) -print(f" Total electronics: {count_result['count']}") +print(f" Total electronics: {count_result.count}") # Query 7: Get only keys and scores (no content) print("\n7. Get keys only (select={}):") @@ -142,16 +144,16 @@ limit=2, ) for result in results: - print(f" Key: {result['key']}, Score: {result['score']}") - print(f" Has data field: {'data' in result}") + print(f" Key: {result.key}, Score: {result.score}") + print(f" Has data field: {result.data is not None}") # Describe the index print("\n8. Index description:") description = products_index.describe() -print(f" Name: {description['name']}") -print(f" Type: {description['dataType']}") -print(f" Prefixes: {description['prefixes']}") -print(f" Schema fields: {list(description['schema'].keys())}") +print(f" Name: {description.name}") +print(f" Type: {description.data_type}") +print(f" Prefixes: {description.prefixes}") +print(f" Schema fields: {list(description.schema.keys())}") # ============================================================================= # Example 2: JSON Index with Nested Schema @@ -164,20 +166,16 @@ # Create index with nested schema posts_index = redis.search.create_index( name="posts", - dataType="json", # Data is stored as JSON - prefix="post:", + data_type="json", # Data is stored as JSON + prefixes="post:", schema={ "title": "TEXT", - "author": { - "name": "TEXT", - "email": "TEXT", - }, - "stats": { - "views": {"type": "U64", "fast": True}, - "likes": {"type": "U64", "fast": True}, - }, + "author.name": "TEXT", + "author.email": "TEXT", + "stats.views": {"type": "U64", "fast": True}, + "stats.likes": {"type": "U64", "fast": True}, "published": "BOOL", - } + }, ) # Add blog posts @@ -211,7 +209,7 @@ print("\n1. Posts by author 'John':") results = posts_index.query(filter={"author.name": {"$eq": "John"}}) for result in results: - data = result["data"] + data = result.data print(f" Title: {data['title']}") print(f" Author: {data['author']['name']}") print(f" Views: {data['stats']['views']}") @@ -220,10 +218,10 @@ print("\n2. Posts with more than 1000 views:") results = posts_index.query( filter={"stats.views": {"$gt": 1000}}, - orderBy={"stats.views": "DESC"}, + order_by={"stats.views": "DESC"}, ) for result in results: - data = result["data"] + data = result.data print(f" {data['title']}: {data['stats']['views']} views") # Select nested fields @@ -234,7 +232,7 @@ limit=2, ) for result in results: - data = result["data"] + data = result.data print(f" Email: {data['author']['email']}, Views: {data['stats']['views']}") # ============================================================================= @@ -248,13 +246,13 @@ # Create hash index scores_index = redis.search.create_index( name="scores", - dataType="hash", # Data is stored as Redis hash - prefix="user:", + data_type="hash", # Data is stored as Redis hash + prefixes="user:", schema={ "username": "TEXT", "score": {"type": "U64", "fast": True}, "level": {"type": "U64", "fast": True}, - } + }, ) # Add user data using HSET @@ -273,10 +271,10 @@ print("\n1. Top level 10 players:") results = scores_index.query( filter={"level": {"$eq": 10}}, - orderBy={"score": "DESC"}, + order_by={"score": "DESC"}, ) for result in results: - data = result["data"] + data = result.data print(f" {data['username']}: {data['score']} points") # ============================================================================= @@ -292,22 +290,22 @@ results = products_index.query(filter={"name": {"$fuzzy": "laptopp"}}) print(f" Found {len(results)} results despite typo") for result in results: - print(f" - {result['data']['name']}") + print(f" - {result.data['name']}") # Phrase search print("\n2. Phrase search for 'Wireless Mouse':") results = products_index.query(filter={"name": {"$phrase": "Wireless Mouse"}}) for result in results: - print(f" Found: {result['data']['name']}") + print(f" Found: {result.data['name']}") # Range query print("\n3. Products priced between $10 and $100:") results = products_index.query( filter={"price": {"$gte": 10, "$lte": 100}}, - orderBy={"price": "ASC"}, + order_by={"price": "ASC"}, ) for result in results: - data = result["data"] + data = result.data print(f" {data['name']}: ${data['price']}") # ============================================================================= diff --git a/pyproject.toml b/pyproject.toml index 9550021..f586827 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ name = "upstash-redis" version = "1.5.0" description = "Serverless Redis SDK from Upstash" license = "MIT" -authors = ["Upstash ", "Zgîmbău Tudor "] +authors = ["Upstash "] maintainers = ["Upstash "] readme = "README.md" repository = "https://github.com/upstash/redis-python" @@ -22,6 +22,7 @@ classifiers = [ "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", "Programming Language :: Python :: Implementation :: CPython", "Topic :: Database", "Topic :: Database :: Front-Ends", @@ -38,7 +39,7 @@ pytest = "^8.3.4" pytest-asyncio = "^0.24.0" python-dotenv = "^1.0.1" mypy = "^1.14.1" -ruff = "^0.9.7" +ruff = "^0.14.13" [build-system] requires = ["poetry-core"] diff --git a/tests/commands/asyncio/search/test_search_index_async.py b/tests/commands/asyncio/search/test_search_index_async.py index ac16495..b9a644a 100644 --- a/tests/commands/asyncio/search/test_search_index_async.py +++ b/tests/commands/asyncio/search/test_search_index_async.py @@ -1,19 +1,21 @@ """Async tests for search index functionality.""" +import dataclasses import json import random import string -from typing import Any, Dict, List, TypedDict +from typing import List import pytest import pytest_asyncio from upstash_redis.asyncio import Redis from upstash_redis.commands import AsyncSearchIndexCommands +from upstash_redis.search import Language, DataType, FieldType, Schema -class QueryIndexFixture(TypedDict): - """Type definition for query_index fixture.""" +@dataclasses.dataclass +class QueryIndexFixture: index: AsyncSearchIndexCommands redis: Redis keys: List[str] @@ -44,7 +46,7 @@ async def cleanup_indexes(): try: index = redis.search.index(index_name) await index.drop() - except: + except Exception: pass @@ -62,7 +64,7 @@ async def query_index(self, async_redis_client: Redis, cleanup_indexes): redis = async_redis_client - schema = { + schema: Schema = { "name": "TEXT", "description": "TEXT", "price": {"type": "F64", "fast": True}, @@ -71,7 +73,7 @@ async def query_index(self, async_redis_client: Redis, cleanup_indexes): } index = await redis.search.create_index( - name=name, schema=schema, dataType="string", prefix=prefix + name=name, schema=schema, data_type="string", prefixes=prefix ) # Add test data @@ -114,118 +116,99 @@ async def query_index(self, async_redis_client: Redis, cleanup_indexes): # Wait for indexing await index.wait_indexing() - yield {"index": index, "redis": redis, "keys": keys, "name": name} + yield QueryIndexFixture(index=index, redis=redis, keys=keys, name=name) # Cleanup try: await index.drop() - except: + except Exception: pass if keys: await redis.delete(*keys) async def test_query_text_eq(self, query_index: QueryIndexFixture): """Test async querying with $eq on text field.""" - index = query_index["index"] + index = query_index.index result = await index.query(filter={"name": {"$eq": "Laptop"}}) assert len(result) == 2 for r in result: - assert "key" in r - # Verify key has the expected prefix - assert query_index["name"] in r["key"] - # Verify score is numeric - assert "score" in r - assert isinstance(float(r["score"]), float) + assert r.key.startswith(query_index.name) + assert r.score > 0.0 async def test_query_text_fuzzy(self, query_index: QueryIndexFixture): """Test async querying with fuzzy search.""" - index = query_index["index"] + index = query_index.index result = await index.query(filter={"name": {"$fuzzy": "laptopp"}}) # Fuzzy search should find Laptop variants despite typo assert len(result) >= 2 - assert all("key" in r and "score" in r for r in result) async def test_query_text_phrase(self, query_index: QueryIndexFixture): """Test async querying with phrase matching.""" - index = query_index["index"] - result = await index.query(filter={"description": {"$phrase": "wireless mouse"}}) + index = query_index.index + result = await index.query( + filter={"description": {"$phrase": "wireless mouse"}} + ) # Should find the Wireless Mouse item assert len(result) == 1 - assert "key" in result[0] - assert "score" in result[0] - assert isinstance(float(result[0]["score"]), float) - # Check data field exists and has expected fields - if "data" in result[0]: - assert "name" in result[0]["data"] or "description" in result[0]["data"] + assert result[0].data is not None + assert "name" in result[0].data and "description" in result[0].data async def test_query_numeric_gt(self, query_index: QueryIndexFixture): """Test async querying with $gt on numeric field.""" - index = query_index["index"] + index = query_index.index result = await index.query(filter={"price": {"$gt": 500}}) # Should find Laptop Pro (1299.99) and Laptop Basic (599.99) assert len(result) == 2 for item in result: - assert "key" in item - assert "score" in item - assert isinstance(float(item["score"]), float) - # Check data field exists with price field - if "data" in item: - assert "price" in item["data"] - # Verify price is greater than 500 - price_value = item["data"]["price"] - assert isinstance(price_value, (int, float, str)) - assert float(price_value) > 500 + assert item.data is not None + assert "price" in item.data + # Verify price is greater than 500 + price_value = item.data["price"] + assert float(price_value) > 500 async def test_query_numeric_gte(self, query_index: QueryIndexFixture): """Test async querying with $gte on numeric field.""" - index = query_index["index"] + index = query_index.index result = await index.query(filter={"stock": {"$gte": 100}}) # Should find items with stock >= 100 (Laptop Basic, Mouse, Cable) assert len(result) == 3 - assert all("key" in r and "score" in r for r in result) async def test_query_numeric_lt(self, query_index: QueryIndexFixture): """Test async querying with $lt on numeric field.""" - index = query_index["index"] + index = query_index.index result = await index.query(filter={"price": {"$lt": 50}}) # Should find Wireless Mouse (29.99) and USB Cable (9.99) assert len(result) == 2 - assert all("key" in r and "score" in r for r in result) async def test_query_boolean_eq(self, query_index: QueryIndexFixture): """Test async querying with $eq on boolean field.""" - index = query_index["index"] + index = query_index.index result = await index.query(filter={"active": {"$eq": False}}) # Should find USB Cable (active: False) assert len(result) == 1 - assert "key" in result[0] - assert "score" in result[0] - assert isinstance(float(result[0]["score"]), float) - # Check data field has active field - if "data" in result[0]: - assert "active" in result[0]["data"] - # Verify active is False - assert result[0]["data"]["active"] in (False, "false", "False", 0, "0") + assert result[0].data is not None + assert "active" in result[0].data + # Verify active is False + assert result[0].data["active"] is False async def test_query_with_limit(self, query_index: QueryIndexFixture): """Test async querying with limit option.""" - index = query_index["index"] + index = query_index.index result = await index.query(filter={"active": {"$eq": True}}, limit=2) # Limit should restrict results to 2 even though 3 items match assert len(result) == 2 - assert all("key" in r and "score" in r for r in result) async def test_query_with_pagination(self, query_index: QueryIndexFixture): """Test async querying with offset for pagination.""" - index = query_index["index"] + index = query_index.index first_page = await index.query( filter={"active": {"$eq": True}}, limit=2, offset=0 ) @@ -238,100 +221,81 @@ async def test_query_with_pagination(self, query_index: QueryIndexFixture): # Should get 1 item on second page (3 total active items) assert len(second_page) == 1 # Pages should have different items - first_keys = {r["key"] for r in first_page} - second_keys = {r["key"] for r in second_page} + first_keys = {r.key for r in first_page} + second_keys = {r.key for r in second_page} assert first_keys.isdisjoint(second_keys) async def test_query_with_sort_asc(self, query_index: QueryIndexFixture): """Test async querying with sortBy ascending.""" - index = query_index["index"] + index = query_index.index result = await index.query( filter={"active": {"$eq": True}}, - orderBy={"price": "ASC"}, + order_by={"price": "ASC"}, limit=3, ) # Should get 3 active items sorted by price ascending assert len(result) == 3 - assert all("key" in r and "score" in r for r in result) async def test_query_with_sort_desc(self, query_index: QueryIndexFixture): """Test async querying with sortBy descending.""" - index = query_index["index"] + index = query_index.index result = await index.query( filter={"active": {"$eq": True}}, - orderBy={"price": "DESC"}, + order_by={"price": "DESC"}, limit=3, ) # Should get 3 active items sorted by price descending assert len(result) == 3 - assert all("key" in r and "score" in r for r in result) async def test_query_with_all_fields(self, query_index: QueryIndexFixture): """Test querying returns all fields when select is not specified.""" - index = query_index["index"] + index = query_index.index result = await index.query(filter={"name": {"$eq": "Wireless Mouse"}}, limit=1) # Should get complete result with all fields assert len(result) == 1 item = result[0] - + # Verify structure - assert "key" in item - assert "score" in item - assert isinstance(float(item["score"]), float) - assert "data" in item - # Verify all fields are present in data (values are strings for string dataType) - data = item["data"] + data = item.data + assert data is not None assert data["name"] == "Wireless Mouse" assert data["description"] == "Ergonomic wireless mouse" - assert float(data["price"]) == 29.99 - assert int(data["stock"]) == 200 - assert data["active"] in (True, "true", "True") + assert data["price"] == 29.99 + assert data["stock"] == 200 + assert data["active"] is True async def test_query_with_no_content(self, query_index: QueryIndexFixture): """Test querying with select={{}} returns only keys and scores.""" - index = query_index["index"] + index = query_index.index result = await index.query( - filter={"name": {"$eq": "Wireless Mouse"}}, - select={}, - limit=1 + filter={"name": {"$eq": "Wireless Mouse"}}, select={}, limit=1 ) # Should get only key and score, no data assert len(result) == 1 - item = result[0] - - assert "key" in item - assert "score" in item - assert isinstance(float(item["score"]), float) # With noContent, data field should not be present - assert "data" not in item + assert result[0].data is None async def test_query_with_specific_field(self, query_index: QueryIndexFixture): """Test querying with select={{"price": True}} returns only specified field.""" - index = query_index["index"] + index = query_index.index result = await index.query( - filter={"name": {"$eq": "Wireless Mouse"}}, - select={"price": True}, - limit=1 + filter={"name": {"$eq": "Wireless Mouse"}}, select={"price": True}, limit=1 ) # Should get key, score, and only the price field assert len(result) == 1 item = result[0] - - assert "key" in item - assert "score" in item - assert isinstance(float(item["score"]), float) - assert "data" in item - + # Verify only price field is present (value is string for string dataType) - data = item["data"] + data = item.data + assert data is not None assert "price" in data - assert float(data["price"]) == 29.99 + assert data["price"] == "29.99" # Other fields should not be present assert "name" not in data assert "description" not in data @@ -354,13 +318,13 @@ async def count_index(self, async_redis_client, cleanup_indexes): redis = async_redis_client schema = { - "category": {"type": "TEXT", "noTokenize": True}, + "category": {"type": "TEXT", "no_tokenize": True}, "price": {"type": "F64", "fast": True}, "active": "BOOL", } index = await redis.search.create_index( - name=name, schema=schema, dataType="string", prefix=prefix + name=name, schema=schema, data_type="string", prefixes=prefix ) # Add test data @@ -379,31 +343,29 @@ async def count_index(self, async_redis_client, cleanup_indexes): # Wait for indexing await index.wait_indexing() - yield {"index": index, "redis": redis, "keys": keys, "name": name} + yield QueryIndexFixture(index=index, redis=redis, keys=keys, name=name) # Cleanup try: await index.drop() - except: + except Exception: pass if keys: await redis.delete(*keys) async def test_count_matching_documents(self, count_index): """Test async counting matching documents.""" - index = count_index["index"] + index = count_index.index result = await index.count(filter={"category": {"$eq": "electronics"}}) - assert "count" in result - assert result["count"] >= 2 + assert result.count >= 2 async def test_count_with_numeric_filter(self, count_index): """Test async counting with numeric filter.""" - index = count_index["index"] + index = count_index.index result = await index.count(filter={"price": {"$gt": 100}}) - assert "count" in result - assert result["count"] >= 2 + assert result.count >= 2 @pytest.mark.asyncio @@ -417,31 +379,27 @@ async def test_describe_index(self, async_redis_client: Redis, cleanup_indexes): redis = async_redis_client - schema = { + schema: Schema = { "name": "TEXT", "price": {"type": "F64", "fast": True}, "active": "BOOL", } index = await redis.search.create_index( - name=name, schema=schema, dataType="string", prefix=f"{name}:" + name=name, schema=schema, data_type="string", prefixes=f"{name}:" ) description = await index.describe() assert description is not None - assert "name" in description - assert description["name"] == name - assert "schema" in description - assert "name" in description["schema"] - assert "price" in description["schema"] - assert "active" in description["schema"] + assert description.name == name + assert description.language == Language.ENGLISH + assert description.data_type == DataType.STRING - # Cleanup - try: - await index.drop() - except: - pass + assert description.schema["name"].type == FieldType.TEXT + assert description.schema["price"].type == FieldType.F64 + assert description.schema["price"].fast is True + assert description.schema["active"].type == FieldType.BOOL @pytest.mark.asyncio @@ -457,7 +415,7 @@ async def test_drop_index(self, async_redis_client): schema = {"name": "TEXT"} index = await redis.search.create_index( - name=name, schema=schema, dataType="string", prefix=f"{name}:" + name=name, schema=schema, data_type="string", prefixes=f"{name}:" ) # Drop the index diff --git a/tests/commands/search/test_search_advanced.py b/tests/commands/search/test_search_advanced.py index 2fc407b..3eadfa1 100644 --- a/tests/commands/search/test_search_advanced.py +++ b/tests/commands/search/test_search_advanced.py @@ -1,7 +1,8 @@ """Tests for JSON and nested schema search functionality.""" +import dataclasses import json -from typing import List, TypedDict, Generator +from typing import List, Generator import pytest @@ -9,8 +10,8 @@ from upstash_redis.commands import SearchIndexCommands -class IndexFixture(TypedDict): - """Type definition for index fixtures.""" +@dataclasses.dataclass +class IndexFixture: index: SearchIndexCommands redis: Redis keys: List[str] @@ -46,14 +47,14 @@ def json_index(self) -> Generator[IndexFixture, None, None]: schema = { "name": "TEXT", "description": "TEXT", - "category": {"type": "TEXT", "noTokenize": True}, + "category": {"type": "TEXT", "no_tokenize": True}, "price": {"type": "F64", "fast": True}, "stock": {"type": "U64", "fast": True}, "active": "BOOL", } index = redis.search.create_index( - name=name, schema=schema, dataType="json", prefix=prefix + name=name, schema=schema, data_type="json", prefixes=prefix ) # Add test data @@ -99,57 +100,57 @@ def json_index(self) -> Generator[IndexFixture, None, None]: index.wait_indexing() - yield {"index": index, "redis": redis, "keys": keys, "name": name} + yield IndexFixture(index=index, redis=redis, keys=keys, name=name) # Cleanup try: index.drop() - except: + except Exception: pass if keys: redis.delete(*keys) def test_query_json_text_eq(self, json_index: IndexFixture): """Test querying JSON index with $eq on text field.""" - index = json_index["index"] + index = json_index.index result = index.query(filter={"name": {"$eq": "Laptop"}}) assert len(result) > 0 def test_query_json_fuzzy(self, json_index: IndexFixture): """Test querying JSON index with fuzzy search.""" - index = json_index["index"] + index = json_index.index result = index.query(filter={"name": {"$fuzzy": "laptopp"}}) assert len(result) > 0 def test_query_json_phrase(self, json_index: IndexFixture): """Test querying JSON index with phrase matching.""" - index = json_index["index"] + index = json_index.index result = index.query(filter={"description": {"$phrase": "wireless mouse"}}) assert len(result) > 0 def test_query_json_regex(self, json_index: IndexFixture): """Test querying JSON index with regex pattern.""" - index = json_index["index"] + index = json_index.index result = index.query(filter={"name": {"$regex": "Laptop.*"}}) assert len(result) > 0 def test_query_json_numeric_gt(self, json_index: IndexFixture): """Test querying JSON index with $gt on numeric field.""" - index = json_index["index"] + index = json_index.index result = index.query(filter={"price": {"$gt": 500}}) assert len(result) > 0 def test_query_json_with_sorting(self, json_index: IndexFixture): """Test querying JSON index with sorting.""" - index = json_index["index"] + index = json_index.index result = index.query( filter={"category": {"$eq": "electronics"}}, - orderBy={"price": "DESC"}, + order_by={"price": "DESC"}, limit=3, ) @@ -171,7 +172,7 @@ def hash_index(self) -> Generator[IndexFixture, None, None]: schema = {"name": "TEXT", "score": {"type": "U64", "fast": True}} index = redis.search.create_index( - name=name, schema=schema, dataType="hash", prefix=prefix + name=name, schema=schema, data_type="hash", prefixes=prefix ) # Add test data using HSET @@ -188,27 +189,27 @@ def hash_index(self) -> Generator[IndexFixture, None, None]: index.wait_indexing() - yield {"index": index, "redis": redis, "keys": keys, "name": name} + yield IndexFixture(index=index, redis=redis, keys=keys, name=name) # Cleanup try: index.drop() - except: + except Exception: pass if keys: redis.delete(*keys) def test_query_hash_by_text(self, hash_index: IndexFixture): """Test querying hash index by text field.""" - index = hash_index["index"] + index = hash_index.index result = index.query(filter={"name": {"$eq": "Alice"}}) assert len(result) > 0 def test_query_hash_with_sorting(self, hash_index: IndexFixture): """Test querying hash index with sorting.""" - index = hash_index["index"] - result = index.query(filter={"score": {"$gte": 80}}, orderBy={"score": "DESC"}) + index = hash_index.index + result = index.query(filter={"score": {"$gte": 80}}, order_by={"score": "DESC"}) assert len(result) > 0 @@ -227,12 +228,14 @@ def nested_string_index(self) -> Generator[IndexFixture, None, None]: schema = { "title": "TEXT", - "author": {"name": "TEXT", "email": "TEXT"}, - "stats": {"views": {"type": "U64", "fast": True}, "likes": {"type": "U64", "fast": True}}, + "author.name": "TEXT", + "author.email": "TEXT", + "stats.views": {"type": "U64", "fast": True}, + "stats.likes": {"type": "U64", "fast": True}, } index = redis.search.create_index( - name=name, schema=schema, dataType="string", prefix=prefix + name=name, schema=schema, data_type="string", prefixes=prefix ) test_data = [ @@ -260,37 +263,37 @@ def nested_string_index(self) -> Generator[IndexFixture, None, None]: index.wait_indexing() - yield {"index": index, "redis": redis, "keys": keys, "name": name} + yield IndexFixture(index=index, redis=redis, keys=keys, name=name) # Cleanup try: index.drop() - except: + except Exception: pass if keys: redis.delete(*keys) def test_query_nested_text_field(self, nested_string_index: IndexFixture): """Test querying nested text field.""" - index = nested_string_index["index"] + index = nested_string_index.index result = index.query(filter={"author.name": {"$eq": "John"}}) assert len(result) > 0 def test_query_nested_numeric_field(self, nested_string_index: IndexFixture): """Test querying nested numeric field.""" - index = nested_string_index["index"] + index = nested_string_index.index result = index.query(filter={"stats.views": {"$eq": 1000}}) assert len(result) > 0 def test_query_nested_with_sorting(self, nested_string_index: IndexFixture): """Test querying with sorting on nested field.""" - index = nested_string_index["index"] + index = nested_string_index.index result = index.query( filter={"author.name": {"$eq": "John"}}, select={"author.email": True}, - orderBy={"stats.views": "DESC"}, + order_by={"stats.views": "DESC"}, ) assert len(result) > 0 @@ -310,12 +313,14 @@ def nested_json_index(self) -> Generator[IndexFixture, None, None]: schema = { "title": "TEXT", - "author": {"name": "TEXT", "email": "TEXT"}, - "stats": {"views": {"type": "U64", "fast": True}, "likes": {"type": "U64", "fast": True}}, + "author.name": "TEXT", + "author.email": "TEXT", + "stats.views": {"type": "U64", "fast": True}, + "stats.likes": {"type": "U64", "fast": True}, } index = redis.search.create_index( - name=name, schema=schema, dataType="json", prefix=prefix + name=name, schema=schema, data_type="json", prefixes=prefix ) test_data = [ @@ -343,37 +348,37 @@ def nested_json_index(self) -> Generator[IndexFixture, None, None]: index.wait_indexing() - yield {"index": index, "redis": redis, "keys": keys, "name": name} + yield IndexFixture(index=index, redis=redis, keys=keys, name=name) # Cleanup try: index.drop() - except: + except Exception: pass if keys: redis.delete(*keys) def test_query_nested_json_text_field(self, nested_json_index: IndexFixture): """Test querying nested text field in JSON index.""" - index = nested_json_index["index"] + index = nested_json_index.index result = index.query(filter={"author.name": {"$eq": "John"}}) assert len(result) > 0 def test_query_nested_json_numeric_field(self, nested_json_index: IndexFixture): """Test querying nested numeric field in JSON index.""" - index = nested_json_index["index"] + index = nested_json_index.index result = index.query(filter={"stats.views": {"$eq": 1000}}) assert len(result) > 0 def test_query_nested_json_with_sorting(self, nested_json_index: IndexFixture): """Test querying with sorting on nested field in JSON index.""" - index = nested_json_index["index"] + index = nested_json_index.index result = index.query( filter={"author.name": {"$eq": "John"}}, select={"author.email": True}, - orderBy={"stats.views": "DESC"}, + order_by={"stats.views": "DESC"}, ) assert len(result) > 0 diff --git a/tests/commands/search/test_search_index.py b/tests/commands/search/test_search_index.py index 1caa854..955d8f6 100644 --- a/tests/commands/search/test_search_index.py +++ b/tests/commands/search/test_search_index.py @@ -1,25 +1,26 @@ """Tests for search index functionality.""" +import dataclasses import json -import time -from typing import Dict, Any, List, TypedDict, Generator +from typing import List, Generator import pytest from upstash_redis import Redis from upstash_redis.commands import SearchIndexCommands +from upstash_redis.search import Language, DataType, FieldType, Schema -class StringIndexFixture(TypedDict): - """Type definition for string_index fixture.""" +@dataclasses.dataclass +class StringIndexFixture: index: SearchIndexCommands redis: Redis keys: List[str] name: str -class CountIndexFixture(TypedDict): - """Type definition for count_index fixture.""" +@dataclasses.dataclass() +class CountIndexFixture: index: SearchIndexCommands redis: Redis keys: List[str] @@ -53,39 +54,43 @@ def cleanup_indexes(): try: index = redis.search.index(index_name) index.drop() - except: + except Exception: pass class TestCreateIndex: """Tests for creating search indexes.""" - def test_create_string_index_simple_schema(self, redis_client: Redis, cleanup_indexes): + def test_create_string_index_simple_schema( + self, redis_client: Redis, cleanup_indexes + ): """Test creating a string index with simple schema.""" name = f"test-string-{random_id()}" cleanup_indexes.append(name) - schema = {"name": "TEXT", "age": {"type": "U64", "fast": True}} + schema: Schema = {"name": "TEXT", "age": {"type": "U64", "fast": True}} index = redis_client.search.create_index( - name=name, schema=schema, dataType="string", prefix=f"{name}:" + name=name, schema=schema, data_type="string", prefixes=f"{name}:" ) assert index is not None - assert index.name == name - def test_create_string_index_nested_schema(self, redis_client: Redis, cleanup_indexes): + def test_create_string_index_nested_schema( + self, redis_client: Redis, cleanup_indexes + ): """Test creating a string index with nested schema.""" name = f"test-string-nested-{random_id()}" cleanup_indexes.append(name) - schema = { + schema: Schema = { "title": "TEXT", - "metadata": {"author": "TEXT", "views": {"type": "U64", "fast": True}}, + "metadata.author": "TEXT", + "metadata.views": {"type": "U64", "fast": True}, } index = redis_client.search.create_index( - name=name, schema=schema, dataType="string", prefix=f"{name}:" + name=name, schema=schema, data_type="string", prefixes=f"{name}:" ) assert index is not None @@ -95,40 +100,44 @@ def test_create_hash_index(self, redis_client: Redis, cleanup_indexes): name = f"test-hash-{random_id()}" cleanup_indexes.append(name) - schema = {"title": "TEXT", "count": {"type": "U64", "fast": True}} + schema: Schema = {"title": "TEXT", "count": {"type": "U64", "fast": True}} index = redis_client.search.create_index( - name=name, schema=schema, dataType="hash", prefix=f"{name}:" + name=name, schema=schema, data_type="hash", prefixes=f"{name}:" ) assert index is not None - def test_create_json_index_simple_schema(self, redis_client: Redis, cleanup_indexes): + def test_create_json_index_simple_schema( + self, redis_client: Redis, cleanup_indexes + ): """Test creating a JSON index with simple schema.""" name = f"test-json-{random_id()}" cleanup_indexes.append(name) - schema = {"name": "TEXT", "age": {"type": "U64", "fast": True}} + schema: Schema = {"name": "TEXT", "age": {"type": "U64", "fast": True}} index = redis_client.search.create_index( - name=name, schema=schema, dataType="json", prefix=f"{name}:" + name=name, schema=schema, data_type="json", prefixes=f"{name}:" ) assert index is not None - assert index.name == name - def test_create_json_index_nested_schema(self, redis_client: Redis, cleanup_indexes): + def test_create_json_index_nested_schema( + self, redis_client: Redis, cleanup_indexes + ): """Test creating a JSON index with nested schema.""" name = f"test-json-nested-{random_id()}" cleanup_indexes.append(name) - schema = { + schema: Schema = { "title": "TEXT", - "metadata": {"author": "TEXT", "views": {"type": "U64", "fast": True}}, + "metadata.author": "TEXT", + "metadata.views": {"type": "U64", "fast": True}, } index = redis_client.search.create_index( - name=name, schema=schema, dataType="json", prefix=f"{name}:" + name=name, schema=schema, data_type="json", prefixes=f"{name}:" ) assert index is not None @@ -138,30 +147,32 @@ def test_create_index_with_language(self, redis_client: Redis, cleanup_indexes): name = f"test-lang-{random_id()}" cleanup_indexes.append(name) - schema = {"content": "TEXT"} + schema: Schema = {"content": "TEXT"} index = redis_client.search.create_index( name=name, schema=schema, - dataType="string", - prefix=f"{name}:", + data_type="string", + prefixes=f"{name}:", language="turkish", ) assert index is not None - def test_create_index_with_multiple_prefixes(self, redis_client: Redis, cleanup_indexes): + def test_create_index_with_multiple_prefixes( + self, redis_client: Redis, cleanup_indexes + ): """Test creating an index with multiple prefixes.""" name = f"test-multi-prefix-{random_id()}" cleanup_indexes.append(name) - schema = {"name": "TEXT"} + schema: Schema = {"name": "TEXT"} index = redis_client.search.create_index( name=name, schema=schema, - dataType="hash", - prefix=[f"{name}:users:", f"{name}:profiles:"], + data_type="hash", + prefixes=[f"{name}:users:", f"{name}:profiles:"], ) assert index is not None @@ -182,14 +193,14 @@ def string_index(self) -> Generator[StringIndexFixture, None, None]: schema = { "name": "TEXT", "description": "TEXT", - "category": {"type": "TEXT", "noTokenize": True}, + "category": {"type": "TEXT", "no_tokenize": True}, "price": {"type": "F64", "fast": True}, "stock": {"type": "U64", "fast": True}, "active": "BOOL", } index = redis.search.create_index( - name=name, schema=schema, dataType="string", prefix=prefix + name=name, schema=schema, data_type="string", prefixes=prefix ) # Add test data @@ -244,146 +255,121 @@ def string_index(self) -> Generator[StringIndexFixture, None, None]: # Wait for indexing index.wait_indexing() - yield {"index": index, "redis": redis, "keys": keys, "name": name} + yield StringIndexFixture(index=index, redis=redis, keys=keys, name=name) # Cleanup try: index.drop() - except: + except Exception: pass if keys: redis.delete(*keys) def test_query_text_eq(self, string_index: StringIndexFixture): """Test querying with $eq on text field.""" - index = string_index["index"] + index = string_index.index result = index.query(filter={"name": {"$eq": "Laptop"}}) # Should find 2 laptops: Laptop Pro and Laptop Basic assert len(result) == 2 for r in result: - assert "key" in r - # Verify key has the expected prefix - assert string_index["name"] in r["key"] - # Verify score is numeric - assert "score" in r - assert isinstance(float(r["score"]), float) + assert r.key.startswith(string_index.name) + assert r.score > 0.0 def test_query_text_fuzzy(self, string_index: StringIndexFixture): """Test querying with fuzzy search for typo tolerance.""" - index = string_index["index"] + index = string_index.index result = index.query(filter={"name": {"$fuzzy": "laptopp"}}) # Fuzzy search should find Laptop variants despite typo assert len(result) >= 2 - assert all("key" in r and "score" in r for r in result) + for r in result: + assert r.key.startswith(string_index.name) + assert r.score > 0.0 def test_query_text_phrase(self, string_index: StringIndexFixture): """Test querying with phrase matching.""" - index = string_index["index"] + index = string_index.index result = index.query(filter={"description": {"$phrase": "wireless mouse"}}) # Should find exactly the Wireless Mouse item assert len(result) == 1 - assert "key" in result[0] - assert "score" in result[0] - assert isinstance(float(result[0]["score"]), float) - # Check data field exists - if "data" in result[0]: - # Check description field in data - assert "description" in result[0]["data"] or "name" in result[0]["data"] + # Check description field in data + assert result[0].data is not None + assert "description" in result[0].data and "name" in result[0].data def test_query_text_regex(self, string_index: StringIndexFixture): """Test querying with regex pattern.""" - index = string_index["index"] + index = string_index.index result = index.query(filter={"name": {"$regex": "Laptop.*"}}) # Should find both Laptop Pro and Laptop Basic assert len(result) == 2 - assert all("key" in r and "score" in r for r in result) def test_query_numeric_gt(self, string_index: StringIndexFixture): """Test querying with $gt on numeric field.""" - index = string_index["index"] + index = string_index.index result = index.query(filter={"price": {"$gt": 500}}) # Should find Laptop Pro (1299.99) and Laptop Basic (599.99) assert len(result) == 2 for r in result: - assert "key" in r - assert "score" in r - assert isinstance(float(r["score"]), float) - # Check data field with price - if "data" in r: - assert "price" in r["data"] - # Verify price is greater than 500 - price_value = r["data"]["price"] - assert isinstance(price_value, (int, float, str)) - assert float(price_value) > 500 + assert r.data is not None + assert "price" in r.data + # Verify price is greater than 500 + price_value = r.data["price"] + assert price_value > 500 def test_query_numeric_gte(self, string_index: StringIndexFixture): """Test querying with $gte on numeric field.""" - index = string_index["index"] + index = string_index.index result = index.query(filter={"stock": {"$gte": 100}}) # Should find items with stock >= 100 (Laptop Basic: 100, Mouse: 200, Cable: 500, Case: 300) assert len(result) == 4 - assert all("key" in r and "score" in r for r in result) def test_query_numeric_lt(self, string_index: StringIndexFixture): """Test querying with $lt on numeric field.""" - index = string_index["index"] + index = string_index.index result = index.query(filter={"price": {"$lt": 50}}) # Should find Wireless Mouse (29.99), USB Cable (9.99), Phone Case (19.99) assert len(result) == 3 - assert all("key" in r and "score" in r for r in result) def test_query_numeric_lte(self, string_index: StringIndexFixture): """Test querying with $lte on numeric field.""" - index = string_index["index"] + index = string_index.index result = index.query(filter={"stock": {"$lte": 50}}) # Should find only Laptop Pro (stock: 50) assert len(result) == 1 - assert "key" in result[0] - assert "score" in result[0] - assert isinstance(float(result[0]["score"]), float) - # Check data field with stock - if "data" in result[0]: - assert "stock" in result[0]["data"] - stock_value = result[0]["data"]["stock"] - assert isinstance(stock_value, (int, float, str)) - assert float(stock_value) <= 50 + assert result[0].data is not None + assert "stock" in result[0].data + assert result[0].data["stock"] <= 50 def test_query_boolean_eq(self, string_index: StringIndexFixture): """Test querying with $eq on boolean field.""" - index = string_index["index"] + index = string_index.index result = index.query(filter={"active": {"$eq": False}}) # Should find only Phone Case (active: False) assert len(result) == 1 - assert "key" in result[0] - assert "score" in result[0] - assert isinstance(float(result[0]["score"]), float) - # Check data field with active - if "data" in result[0]: - assert "active" in result[0]["data"] - # Verify active is False - assert result[0]["data"]["active"] in (False, "false", "False", 0, "0") + assert result[0].data is not None + assert "active" in result[0].data + # Verify active is False + assert result[0].data["active"] is False def test_query_with_limit(self, string_index: StringIndexFixture): """Test querying with limit option.""" - index = string_index["index"] + index = string_index.index result = index.query(filter={"category": {"$eq": "electronics"}}, limit=2) # Should limit to exactly 2 results even though 3 match assert len(result) == 2 - assert all("key" in r and "score" in r for r in result) def test_query_with_pagination(self, string_index: StringIndexFixture): """Test querying with offset for pagination.""" - index = string_index["index"] + index = string_index.index first_page = index.query( filter={"category": {"$eq": "electronics"}}, limit=2, offset=0 ) @@ -397,103 +383,95 @@ def test_query_with_pagination(self, string_index: StringIndexFixture): # Should get 1 item on second page (3 total electronics) assert len(second_page) == 1 # Pages should have different items - first_keys = {r["key"] for r in first_page} - second_keys = {r["key"] for r in second_page} + first_keys = {r.key for r in first_page} + second_keys = {r.key for r in second_page} assert first_keys.isdisjoint(second_keys) def test_query_with_sort_asc(self, string_index: StringIndexFixture): """Test querying with sortBy ascending.""" - index = string_index["index"] + index = string_index.index result = index.query( filter={"category": {"$eq": "electronics"}}, - orderBy={"price": "ASC"}, + order_by={"price": "ASC"}, limit=3, ) # Should get all 3 electronics items sorted by price assert len(result) == 3 - assert all("key" in r and "score" in r for r in result) + prev_score = float("-inf") + for r in result: + assert r.score > prev_score + prev_score = r.score def test_query_with_sort_desc(self, string_index: StringIndexFixture): """Test querying with sortBy descending.""" - index = string_index["index"] + index = string_index.index result = index.query( filter={"category": {"$eq": "electronics"}}, - orderBy={"price": "DESC"}, + order_by={"price": "DESC"}, limit=3, ) # Should get all 3 electronics items sorted by price descending assert len(result) == 3 - assert all("key" in r and "score" in r for r in result) + prev_score = float("inf") + for r in result: + assert r.score < prev_score + prev_score = r.score def test_query_with_all_fields(self, string_index: StringIndexFixture): """Test querying returns all fields when select is not specified.""" - index = string_index["index"] + index = string_index.index result = index.query(filter={"name": {"$eq": "Wireless Mouse"}}, limit=1) # Should get complete result with all fields assert len(result) == 1 item = result[0] - - # Verify structure - assert "key" in item - assert "score" in item - assert isinstance(float(item["score"]), float) - assert "data" in item - - # Verify all fields are present in data (values are strings for string dataType) - data = item["data"] + + # Verify all fields are present in data (values are strings for string data_type) + data = item.data + assert data is not None assert data["name"] == "Wireless Mouse" assert data["description"] == "Ergonomic wireless mouse" assert data["category"] == "electronics" - assert float(data["price"]) == 29.99 - assert int(data["stock"]) == 200 - assert data["active"] in (True, "true", "True") + assert data["price"] == 29.99 + assert data["stock"] == 200 + assert data["active"] is True def test_query_no_content(self, string_index: StringIndexFixture): """Test querying with noContent (keys only).""" - index = string_index["index"] + index = string_index.index result = index.query( - filter={"name": {"$eq": "Wireless Mouse"}}, - select={}, - limit=1 + filter={"name": {"$eq": "Wireless Mouse"}}, select={}, limit=1 ) # Should get only key and score, no data field assert len(result) == 1 item = result[0] - - assert "key" in item - assert "score" in item - assert isinstance(float(item["score"]), float) + # With noContent, data field should not be present - assert "data" not in item + assert item.data is None def test_query_with_return_fields(self, string_index: StringIndexFixture): """Test querying with specific return fields.""" - index = string_index["index"] + index = string_index.index result = index.query( filter={"name": {"$eq": "Wireless Mouse"}}, select={"price": True, "stock": True}, - limit=1 + limit=1, ) # Should get only specified fields assert len(result) == 1 item = result[0] - - assert "key" in item - assert "score" in item - assert isinstance(float(item["score"]), float) - assert "data" in item - - # Verify only selected fields are present (values are strings for string dataType) - data = item["data"] + + # Verify only selected fields are present (values are strings for string data_type) + data = item.data + assert data is not None assert "price" in data - assert float(data["price"]) == 29.99 + assert data["price"] == "29.99" assert "stock" in data - assert int(data["stock"]) == 200 + assert data["stock"] == "200" # Other fields should not be present assert "name" not in data assert "description" not in data @@ -516,7 +494,7 @@ def count_index(self) -> Generator[CountIndexFixture, None, None]: schema = {"type": "TEXT", "value": {"type": "U64", "fast": True}} index = redis.search.create_index( - name=name, schema=schema, dataType="string", prefix=prefix + name=name, schema=schema, data_type="string", prefixes=prefix ) # Add test data @@ -527,33 +505,31 @@ def count_index(self) -> Generator[CountIndexFixture, None, None]: index.wait_indexing() - yield {"index": index, "redis": redis, "keys": keys, "name": name} + yield CountIndexFixture(index=index, redis=redis, keys=keys, name=name) # Cleanup try: index.drop() - except: + except Exception: pass if keys: redis.delete(*keys) def test_count_matching_documents(self, count_index: CountIndexFixture): """Test counting matching documents.""" - index = count_index["index"] - result = index.count({"type": {"$eq": "A"}}) + index = count_index.index + result = index.count(filter={"type": {"$eq": "A"}}) - assert "count" in result # Should count exactly 5 documents of type A (indices 0-4) - assert result["count"] == 5 + assert result.count == 5 def test_count_with_numeric_filter(self, count_index: CountIndexFixture): """Test counting with numeric filter.""" - index = count_index["index"] - result = index.count({"value": {"$eq": 5}}) + index = count_index.index + result = index.count(filter={"value": {"$eq": 5}}) - assert "count" in result # Should count exactly 1 document with value 5 - assert result["count"] == 1 + assert result.count == 1 class TestDescribe: @@ -564,20 +540,28 @@ def test_describe_index(self, redis_client: Redis, cleanup_indexes): name = f"test-describe-{random_id()}" cleanup_indexes.append(name) - schema = { - "title": {"type": "TEXT", "noStem": True}, + schema: Schema = { + "title": {"type": "TEXT", "no_stem": True}, "count": {"type": "U64", "fast": True}, "active": "BOOL", } index = redis_client.search.create_index( - name=name, schema=schema, dataType="string", prefix=f"{name}:" + name=name, schema=schema, data_type="string", prefixes=f"{name}:" ) description = index.describe() assert description is not None - assert "name" in description or "schema" in description + assert description.name == name + assert description.language == Language.ENGLISH + assert description.data_type == DataType.STRING + + assert description.schema["active"].type == FieldType.BOOL + assert description.schema["title"].type == FieldType.TEXT + assert description.schema["title"].no_stem is True + assert description.schema["count"].type == FieldType.U64 + assert description.schema["count"].fast is True class TestDrop: @@ -587,15 +571,13 @@ def test_drop_existing_index(self, redis_client: Redis): """Test dropping an existing index.""" name = f"test-drop-{random_id()}" - schema = {"name": "TEXT"} + schema: Schema = {"name": "TEXT"} index = redis_client.search.create_index( - name=name, schema=schema, dataType="string", prefix=f"{name}:" + name=name, schema=schema, data_type="string", prefixes=f"{name}:" ) - result = index.drop() - # Result should be 1 or 0 or "OK" - assert result is not None + index.drop() class TestIndexMethod: @@ -606,4 +588,3 @@ def test_index(self, redis_client: Redis): index = redis_client.search.index("test-index") assert index is not None - assert index.name == "test-index" diff --git a/tests/commands/search/test_search_utils.py b/tests/commands/search/test_search_utils.py deleted file mode 100644 index d1e1f09..0000000 --- a/tests/commands/search/test_search_utils.py +++ /dev/null @@ -1,377 +0,0 @@ -"""Unit tests for search utility functions.""" - -import pytest - -from upstash_redis.search_utils import ( - flatten_schema, - build_create_index_command, - build_query_command, - deserialize_query_response, - deserialize_describe_response, -) - - -class TestFlattenSchema: - """Tests for flatten_schema function.""" - - def test_flatten_simple_schema(self): - """Test flattening a simple schema.""" - schema = {"name": "TEXT", "age": {"type": "U64", "fast": True}, "active": "BOOL"} - - flattened = flatten_schema(schema) - - assert len(flattened) == 3 - assert flattened[0].path == "name" - assert flattened[0].field_type == "TEXT" - assert flattened[1].path == "age" - assert flattened[1].field_type == "U64" - assert flattened[1].fast is True - - def test_flatten_nested_schema(self): - """Test flattening a nested schema.""" - schema = { - "title": "TEXT", - "author": {"name": "TEXT", "email": "TEXT"}, - "stats": {"views": {"type": "U64", "fast": True}}, - } - - flattened = flatten_schema(schema) - - assert len(flattened) == 4 - # Check that nested paths are created correctly - paths = [f.path for f in flattened] - assert "title" in paths - assert "author.name" in paths - assert "author.email" in paths - assert "stats.views" in paths - - def test_flatten_with_field_options(self): - """Test flattening schema with field options.""" - schema = { - "title": {"type": "TEXT", "noStem": True, "noTokenize": True}, - "count": {"type": "U64", "fast": True}, - } - - flattened = flatten_schema(schema) - - title_field = next(f for f in flattened if f.path == "title") - assert title_field.no_stem is True - assert title_field.no_tokenize is True - - -class TestBuildCreateIndexCommand: - """Tests for build_create_index_command function.""" - - def test_build_simple_create_command(self): - """Test building a simple create index command.""" - params = { - "name": "test-idx", - "dataType": "string", - "prefix": "test:", - "schema": {"name": "TEXT", "age": {"type": "U64", "fast": True}}, - } - - command = build_create_index_command(params) - - assert command == [ - "SEARCH.CREATE", - "test-idx", - "ON", - "STRING", - "PREFIX", - "1", - "test:", - "SCHEMA", - "name", - "TEXT", - "age", - "U64", - "FAST", - ] - - def test_build_create_with_language(self): - """Test building create command with language option.""" - params = { - "name": "test-idx", - "dataType": "json", - "prefix": "test:", - "schema": {"content": "TEXT"}, - "language": "turkish", - } - - command = build_create_index_command(params) - - assert command == [ - "SEARCH.CREATE", - "test-idx", - "ON", - "JSON", - "PREFIX", - "1", - "test:", - "LANGUAGE", - "turkish", - "SCHEMA", - "content", - "TEXT", - ] - - def test_build_create_with_multiple_prefixes(self): - """Test building create command with multiple prefixes.""" - params = { - "name": "test-idx", - "dataType": "hash", - "prefix": ["user:", "profile:"], - "schema": {"name": "TEXT"}, - } - - command = build_create_index_command(params) - - assert command == [ - "SEARCH.CREATE", - "test-idx", - "ON", - "HASH", - "PREFIX", - "2", - "user:", - "profile:", - "SCHEMA", - "name", - "TEXT", - ] - - def test_build_create_with_skip_initial_scan(self): - """Test building create command with skipInitialScan.""" - params = { - "name": "test-idx", - "dataType": "string", - "prefix": "test:", - "schema": {"name": "TEXT"}, - "skipInitialScan": True, - } - - command = build_create_index_command(params) - - assert command == [ - "SEARCH.CREATE", - "test-idx", - "SKIPINITIALSCAN", - "ON", - "STRING", - "PREFIX", - "1", - "test:", - "SCHEMA", - "name", - "TEXT", - ] - - def test_build_create_with_field_options(self): - """Test building create command with field options.""" - params = { - "name": "test-idx", - "dataType": "string", - "prefix": "test:", - "schema": { - "title": {"type": "TEXT", "noStem": True, "noTokenize": True}, - "score": {"type": "F64", "fast": True}, - }, - } - - command = build_create_index_command(params) - - assert command == [ - "SEARCH.CREATE", - "test-idx", - "ON", - "STRING", - "PREFIX", - "1", - "test:", - "SCHEMA", - "title", - "TEXT", - "NOTOKENIZE", - "NOSTEM", - "score", - "F64", - "FAST", - ] - - -class TestBuildQueryCommand: - """Tests for build_query_command function.""" - - def test_build_simple_query(self): - """Test building a simple query command.""" - command = build_query_command( - "SEARCH.QUERY", "test-idx", {"filter": {"name": {"$eq": "test"}}} - ) - - assert command[0] == "SEARCH.QUERY" - assert command[1] == "test-idx" - assert '{"name":{"$eq":"test"}}' in command[2] or '{"name": {"$eq": "test"}}' in command[2] - - def test_build_query_with_limit(self): - """Test building query with limit.""" - command = build_query_command( - "SEARCH.QUERY", "test-idx", {"filter": {"name": {"$eq": "test"}}, "limit": 10} - ) - - assert command == [ - "SEARCH.QUERY", - "test-idx", - '{"name":{"$eq":"test"}}', - "LIMIT", - "10", - ] - - def test_build_query_with_offset(self): - """Test building query with offset.""" - command = build_query_command( - "SEARCH.QUERY", "test-idx", {"filter": {"name": {"$eq": "test"}}, "limit": 10, "offset": 5} - ) - - assert command == [ - "SEARCH.QUERY", - "test-idx", - '{"name":{"$eq":"test"}}', - "LIMIT", - "10", - "OFFSET", - "5", - ] - - def test_build_query_with_sorting(self): - """Test building query with sorting.""" - command = build_query_command( - "SEARCH.QUERY", - "test-idx", - {"filter": {"name": {"$eq": "test"}}, "orderBy": {"score": "DESC"}}, - ) - - assert command == [ - "SEARCH.QUERY", - "test-idx", - '{"name":{"$eq":"test"}}', - "SORTBY", - "score", - "DESC", - ] - - def test_build_query_with_select_nocontent(self): - """Test building query with noContent.""" - command = build_query_command( - "SEARCH.QUERY", "test-idx", {"filter": {"name": {"$eq": "test"}}, "select": {}} - ) - - assert command == [ - "SEARCH.QUERY", - "test-idx", - '{"name":{"$eq":"test"}}', - "NOCONTENT", - ] - - def test_build_query_with_select_fields(self): - """Test building query with specific return fields.""" - command = build_query_command( - "SEARCH.QUERY", - "test-idx", - {"filter": {"name": {"$eq": "test"}}, "select": {"name": True, "age": True}}, - ) - - assert command == [ - "SEARCH.QUERY", - "test-idx", - '{"name":{"$eq":"test"}}', - "RETURN", - "2", - "name", - "age", - ] - - def test_build_query_numeric_filters(self): - """Test building query with numeric filters.""" - # Test JSON format - command = build_query_command("SEARCH.QUERY", "test-idx", {"filter": {"price": {"$gt": 100}}}) - assert command == [ - "SEARCH.QUERY", - "test-idx", - '{"price":{"$gt":100}}', - ] - - def test_build_query_text_filters(self): - """Test building query with text filters.""" - # Test JSON format - command = build_query_command( - "SEARCH.QUERY", "test-idx", {"filter": {"name": {"$fuzzy": "test"}}} - ) - assert command == [ - "SEARCH.QUERY", - "test-idx", - '{"name":{"$fuzzy":"test"}}', - ] - - -class TestDeserializeQueryResponse: - """Tests for deserialize_query_response function.""" - - def test_deserialize_simple_response(self): - """Test deserializing a simple query response.""" - raw = [["key1", "0.5", [["name", "test"], ["age", 25]]]] - - result = deserialize_query_response(raw) - - assert len(result) == 1 - assert result[0]["key"] == "key1" - assert result[0]["score"] == "0.5" - assert "data" in result[0] - - def test_deserialize_nocontent_response(self): - """Test deserializing a response with no content.""" - raw = [["key1", "0.5"]] - - result = deserialize_query_response(raw) - - assert len(result) == 1 - assert result[0]["key"] == "key1" - assert result[0]["score"] == "0.5" - assert "data" not in result[0] - - def test_deserialize_nested_fields(self): - """Test deserializing response with nested fields.""" - raw = [["key1", "0.5", [["author.name", "John"], ["author.email", "john@example.com"]]]] - - result = deserialize_query_response(raw) - - assert len(result) == 1 - assert "data" in result[0] - # The deserialization should create nested structure - # (implementation may vary) - - -class TestDeserializeDescribeResponse: - """Tests for deserialize_describe_response function.""" - - def test_deserialize_describe(self): - """Test deserializing a describe response.""" - raw = [ - "name", - "test-idx", - "type", - "STRING", - "prefixes", - ["test:"], - "schema", - [["name", "TEXT"], ["age", "U64", "FAST"]], - ] - - result = deserialize_describe_response(raw) - - assert result["name"] == "test-idx" - assert result["dataType"] == "string" - assert result["prefixes"] == ["test:"] - assert "schema" in result - assert "name" in result["schema"] - assert result["schema"]["name"]["type"] == "TEXT" diff --git a/tests/test_formatters.py b/tests/test_formatters.py index c6fb18b..6da1c39 100644 --- a/tests/test_formatters.py +++ b/tests/test_formatters.py @@ -11,6 +11,7 @@ format_search_describe_response, format_search_count_response, ) +from upstash_redis.search import CountResult, FieldType, Language, DataType from upstash_redis.utils import GeoSearchResult @@ -182,16 +183,16 @@ def test_format_search_query_response() -> None: ["doc:2", 0.85, [["name", "Phone"], ["price", "699"]]], ] result = format_search_query_response(raw_response, None, None) - + assert len(result) == 2 - assert result[0]["key"] == "doc:1" - assert result[0]["score"] == 0.95 - assert result[0]["data"]["name"] == "Laptop" - assert result[0]["data"]["price"] == "999" - assert result[1]["key"] == "doc:2" - assert result[1]["score"] == 0.85 - assert result[1]["data"]["name"] == "Phone" - assert result[1]["data"]["price"] == "699" + assert result[0].key == "doc:1" + assert result[0].score == 0.95 + assert result[0].data["name"] == "Laptop" + assert result[0].data["price"] == "999" + assert result[1].key == "doc:2" + assert result[1].score == 0.85 + assert result[1].data["name"] == "Phone" + assert result[1].data["price"] == "699" def test_format_search_query_response_with_nested_paths() -> None: @@ -200,26 +201,26 @@ def test_format_search_query_response_with_nested_paths() -> None: ["doc:1", 0.95, [["user.name", "John"], ["user.age", "30"]]], ] result = format_search_query_response(raw_response, None, None) - + assert len(result) == 1 - assert result[0]["key"] == "doc:1" - assert result[0]["score"] == 0.95 - assert result[0]["data"]["user"]["name"] == "John" - assert result[0]["data"]["user"]["age"] == "30" + assert result[0].key == "doc:1" + assert result[0].score == 0.95 + assert result[0].data["user"]["name"] == "John" + assert result[0].data["user"]["age"] == "30" def test_format_search_query_response_with_dollar_key() -> None: # Test query response with $ key (full document) raw_response = [ - ["doc:1", 1.0, [["$", {"name": "Laptop", "price": 999}]]], + ["doc:1", 1.0, [["$", '{"name": "Laptop", "price": 999}']]], ] result = format_search_query_response(raw_response, None, None) - + assert len(result) == 1 - assert result[0]["key"] == "doc:1" - assert result[0]["score"] == 1.0 - assert result[0]["data"]["name"] == "Laptop" - assert result[0]["data"]["price"] == 999 + assert result[0].key == "doc:1" + assert result[0].score == 1.0 + assert result[0].data["name"] == "Laptop" + assert result[0].data["price"] == 999 def test_format_search_query_response_without_fields() -> None: @@ -229,69 +230,74 @@ def test_format_search_query_response_without_fields() -> None: ["doc:2", 0.85], ] result = format_search_query_response(raw_response, None, None) - + assert len(result) == 2 - assert result[0]["key"] == "doc:1" - assert result[0]["score"] == 0.95 - assert "data" not in result[0] - assert result[1]["key"] == "doc:2" - assert result[1]["score"] == 0.85 - assert "data" not in result[1] + assert result[0].key == "doc:1" + assert result[0].score == 0.95 + assert result[0].data is None + assert result[1].key == "doc:2" + assert result[1].score == 0.85 + assert result[1].data is None def test_format_search_describe_response() -> None: # Test describe response raw_response = [ - "name", "myindex", - "type", "JSON", - "prefixes", ["product:", "item:"], - "language", "english", - "schema", [ + "name", + "myindex", + "type", + "JSON", + "prefixes", + ["product:", "item:"], + "language", + "english", + "schema", + [ ["name", "TEXT", "NOSTEM"], ["price", "F64", "FAST"], ["active", "BOOL"], ], ] result = format_search_describe_response(raw_response, None, None) - - assert result["name"] == "myindex" - assert result["dataType"] == "json" - assert result["prefixes"] == ["product:", "item:"] - assert result["language"] == "english" - assert "name" in result["schema"] - assert result["schema"]["name"]["type"] == "TEXT" - assert result["schema"]["name"]["noStem"] is True - assert "price" in result["schema"] - assert result["schema"]["price"]["type"] == "F64" - assert result["schema"]["price"]["fast"] is True - assert "active" in result["schema"] - assert result["schema"]["active"]["type"] == "BOOL" + + assert result.name == "myindex" + assert result.data_type == DataType.JSON + assert result.prefixes == ["product:", "item:"] + assert result.language == Language.ENGLISH + assert "name" in result.schema + assert result.schema["name"].type == FieldType.TEXT + assert result.schema["name"].no_stem is True + assert "price" in result.schema + assert result.schema["price"].type == FieldType.F64 + assert result.schema["price"].fast is True + assert "active" in result.schema + assert result.schema["active"].type == FieldType.BOOL def test_format_search_describe_response_with_all_options() -> None: # Test describe response with all field options raw_response = [ - "name", "fullindex", - "type", "HASH", - "schema", [ + "name", + "fullindex", + "type", + "HASH", + "schema", + [ ["title", "TEXT", "NOSTEM", "NOTOKENIZE"], ["score", "I64", "FAST"], ], ] result = format_search_describe_response(raw_response, None, None) - - assert result["name"] == "fullindex" - assert result["dataType"] == "hash" - assert result["schema"]["title"]["type"] == "TEXT" - assert result["schema"]["title"]["noStem"] is True - assert result["schema"]["title"]["noTokenize"] is True - assert result["schema"]["score"]["type"] == "I64" - assert result["schema"]["score"]["fast"] is True + + assert result.name == "fullindex" + assert result.data_type == DataType.HASH + assert result.schema["title"].type == FieldType.TEXT + assert result.schema["title"].no_stem is True + assert result.schema["title"].no_tokenize is True + assert result.schema["score"].type == FieldType.I64 + assert result.schema["score"].fast is True def test_format_search_count_response() -> None: # Test count response with int - assert format_search_count_response(42, None, None) == {"count": 42} - - # Test count response with string - assert format_search_count_response("100", None, None) == {"count": 100} + assert format_search_count_response(42, None, None) == CountResult(count=42) diff --git a/upstash_redis/asyncio/client.py b/upstash_redis/asyncio/client.py index 2af7bd1..6a8fdfa 100644 --- a/upstash_redis/asyncio/client.py +++ b/upstash_redis/asyncio/client.py @@ -154,7 +154,7 @@ def pipeline(self) -> "AsyncPipeline": http=self._http, multi_exec="pipeline", set_sync_token_header_fn=self._maybe_set_sync_token_header, - original_client=self + original_client=self, ) def multi(self) -> "AsyncPipeline": @@ -167,7 +167,7 @@ def multi(self) -> "AsyncPipeline": http=self._http, multi_exec="multi-exec", set_sync_token_header_fn=self._maybe_set_sync_token_header, - original_client=self + original_client=self, ) diff --git a/upstash_redis/client.py b/upstash_redis/client.py index ac88d64..75ea999 100644 --- a/upstash_redis/client.py +++ b/upstash_redis/client.py @@ -156,7 +156,7 @@ def pipeline(self) -> "Pipeline": http=self._http, multi_exec="pipeline", set_sync_token_header_fn=self._maybe_set_sync_token_header, - original_client=self + original_client=self, ) def multi(self) -> "Pipeline": @@ -169,7 +169,7 @@ def multi(self) -> "Pipeline": http=self._http, multi_exec="multi-exec", set_sync_token_header_fn=self._maybe_set_sync_token_header, - original_client=self + original_client=self, ) diff --git a/upstash_redis/commands.py b/upstash_redis/commands.py index b2d3ac8..a937d64 100644 --- a/upstash_redis/commands.py +++ b/upstash_redis/commands.py @@ -1,7 +1,26 @@ import datetime -from typing import Any, Awaitable, Dict, List, Literal, Mapping, Optional, Tuple, Union +import json +from typing import ( + Any, + Awaitable, + Dict, + List, + Literal, + Mapping, + Optional, + Tuple, + Union, +) -from upstash_redis.typing import FloatMinMaxT, ValueT, JSONValueT +from upstash_redis.search import ( + DataType, + Language, + Schema, + FieldType, + Order, + HighlightOptions, +) +from upstash_redis.typing import FloatMinMaxT, JSONValueT, ValueT from upstash_redis.utils import ( handle_georadius_write_exceptions, handle_geosearch_exceptions, @@ -9,8 +28,6 @@ handle_zrangebylex_exceptions, number_are_not_none, ) -from upstash_redis.search_types import CreateIndexParams, QueryOptions -from upstash_redis.search_utils import build_create_index_command, build_query_command ResponseT = Union[Awaitable, Any] @@ -5489,41 +5506,71 @@ def type(self, key: str, path: str = "$") -> ResponseT: class SearchCommands: """ Search commands namespace. - + This class provides methods for search index operations. - It follows the same pattern as JsonCommands. """ + def __init__(self, client: Commands): self.client = client - + def create_index( self, *, name: str, - schema, - dataType: str, - prefix, - language = None, - skipInitialScan: bool = False, - existsOk: bool = False, + schema: Schema, + data_type: Union[DataType, str], + prefixes: Union[str, List[str]], + language: Optional[Union[Language, str]] = None, + skip_initial_scan: bool = False, + exists_ok: bool = False, ) -> ResponseT: """Create a new search index.""" - params: CreateIndexParams = { - "name": name, - "schema": schema, - "dataType": dataType, - "prefix": prefix, - } - if language is not None: - params["language"] = language - if skipInitialScan: - params["skipInitialScan"] = skipInitialScan - if existsOk: - params["existsOk"] = existsOk - - command = build_create_index_command(params) + + if isinstance(prefixes, str): + prefixes = [prefixes] + + command: List = [ + "SEARCH.CREATE", + name, + "ON", + data_type, + "PREFIX", + len(prefixes), + *prefixes, + ] + + if language: + command.extend(("LANGUAGE", language)) + + if skip_initial_scan: + command.append("SKIPINITIALSCAN") + + if exists_ok: + command.append("EXISTSOK") + + command.append("SCHEMA") + + for path, value in schema.items(): + if isinstance(value, (str, FieldType)): + command.extend((path, value)) + else: + field_type = value["type"] + command.extend((path, field_type)) + + if "alias" in value: + command.extend(("AS", value["alias"])) + + if "fast" in value and value["fast"]: + command.append("FAST") + + if "no_stem" in value and value["no_stem"]: + command.append("NOSTEM") + + if "no_tokenize" in value and value["no_tokenize"]: + command.append("NOTOKENIZE") + return self.client.execute(command) - + def index(self, name: str) -> "SearchIndexCommands": """Initialize a SearchIndexCommands instance for an existing index.""" return SearchIndexCommands(self.client, name) @@ -5532,88 +5579,105 @@ def index(self, name: str) -> "SearchIndexCommands": class SearchIndexCommands: """ Commands for interacting with a specific search index. - + This class provides methods to query, count, describe, and drop an index. """ + def __init__(self, client: Commands, name: str): self.client = client self.name = name - - def wait_indexing(self) -> ResponseT: + + def query( + self, + *, + filter: Dict[str, Any], + limit: Optional[int] = None, + offset: Optional[int] = None, + order_by: Optional[Dict[str, Union[Order, str]]] = None, + select: Optional[Dict[str, bool]] = None, + highlight: Optional[HighlightOptions] = None, + ) -> ResponseT: """ - Wait for the index to finish indexing all existing keys. - + Query the index with filters and options. + Example: ```python - index.wait_indexing() + results = index.query(filter={"name": {"$eq": "Laptop"}}) ``` """ - command = ["SEARCH.WAITINDEXING", self.name] + command: List = ["SEARCH.QUERY", self.name, json.dumps(filter)] + + if limit: + command.extend(("LIMIT", limit)) + + if offset: + command.extend(("OFFSET", offset)) + + if order_by: + for field, order in order_by.items(): + command.extend(("ORDERBY", field, order)) + break + + if select is not None: + if not select: + command.append("NOCONTENT") + else: + selected_fields = [] + for field, selected in select.items(): + if selected: + selected_fields.append(field) + + command.extend(("SELECT", len(selected_fields), *selected_fields)) + + if highlight: + command.extend(("HIGHLIGHT", "FIELDS", highlight["fields"])) + if highlight["tags"]: + open_tag, close_tag = highlight["tags"] + command.extend(("TAGS", open_tag, close_tag)) + return self.client.execute(command) - - def describe(self) -> ResponseT: + + def count(self, *, filter: Dict[str, Any]) -> ResponseT: """ - Get detailed information about the index structure. - + Count documents matching a filter. + Example: ```python - description = index.describe() + result = index.count({"active": {"$eq": True}}) ``` """ - command = ["SEARCH.DESCRIBE", self.name] + + command: List = ["SEARCH.COUNT", self.name, json.dumps(filter)] return self.client.execute(command) - - def query( - self, - *, - filter, - limit = None, - offset = None, - orderBy = None, - select = None, - highlight = None, - ) -> ResponseT: + + def wait_indexing(self) -> ResponseT: """ - Query the index with filters and options. - + Wait for the index to finish indexing all existing keys. + Example: ```python - results = index.query(filter={"name": {"$eq": "Laptop"}}) + index.wait_indexing() ``` """ - options: QueryOptions = {} - if filter is not None: - options["filter"] = filter - if limit is not None: - options["limit"] = limit - if offset is not None: - options["offset"] = offset - if orderBy is not None: - options["orderBy"] = orderBy - if select is not None: - options["select"] = select - if highlight is not None: - options["highlight"] = highlight - - command = build_query_command("SEARCH.QUERY", self.name, options) + command = ["SEARCH.WAITINDEXING", self.name] return self.client.execute(command) - - def count(self, filter) -> ResponseT: + + def describe(self) -> ResponseT: """ - Count documents matching a filter. - + Get detailed information about the index structure. + Example: ```python - result = index.count({"active": {"$eq": True}}) + description = index.describe() ``` """ - command = build_query_command("SEARCH.COUNT", self.name, {"filter": filter}) + command = ["SEARCH.DESCRIBE", self.name] return self.client.execute(command) - + def drop(self) -> ResponseT: """ Drop (delete) the index. - + Example: ```python result = index.drop() diff --git a/upstash_redis/commands.pyi b/upstash_redis/commands.pyi index aaf4c16..879d679 100644 --- a/upstash_redis/commands.pyi +++ b/upstash_redis/commands.pyi @@ -1,9 +1,18 @@ import datetime from typing import Any, Dict, List, Literal, Mapping, Optional, Tuple, Union +from upstash_redis.search import ( + IndexDescription, + QueryResult, + Schema, + DataType, + Language, + Order, + HighlightOptions, + CountResult, +) from upstash_redis.typing import FloatMinMaxT, ValueT, JSONValueT from upstash_redis.utils import GeoSearchResult -from upstash_redis.search_types import NestedIndexSchema, FlatIndexSchema, IndexDescription, QueryResult class Commands: def bitcount( @@ -2036,22 +2045,21 @@ class SearchCommands: self, *, name: str, - schema: Union[NestedIndexSchema, FlatIndexSchema], - dataType: str, - prefix: Union[str, List[str]], - language: Optional[str] = None, - skipInitialScan: bool = False, - existsOk: bool = False, + schema: Schema, + data_type: Union[DataType, str], + prefixes: Union[str, List[str]], + language: Optional[Union[Language, str]] = None, + skip_initial_scan: bool = False, + exists_ok: bool = False, ) -> SearchIndexCommands: ... def index( self, name: str, - schema: Optional[Union[NestedIndexSchema, FlatIndexSchema]] = None ) -> SearchIndexCommands: ... class SearchIndexCommands: - def __init__(self, client: Commands, name: str, schema: Optional[Union[NestedIndexSchema, FlatIndexSchema]] = None): ... - def wait_indexing(self) -> Any: ... + def __init__(self, client: Commands, name: str): ... + def wait_indexing(self) -> None: ... def describe(self) -> IndexDescription: ... def query( self, @@ -2059,12 +2067,12 @@ class SearchIndexCommands: filter: Dict[str, Any], limit: Optional[int] = None, offset: Optional[int] = None, - orderBy: Optional[Dict[str, str]] = None, + order_by: Optional[Dict[str, Union[Order, str]]] = None, select: Optional[Dict[str, bool]] = None, - highlight: Optional[Dict[str, Any]] = None, + highlight: Optional[HighlightOptions] = None, ) -> List[QueryResult]: ... - def count(self, filter: Dict[str, Any]) -> Dict[Literal["count"], int]: ... - def drop(self) -> int: ... + def count(self, *, filter: Dict[str, Any]) -> CountResult: ... + def drop(self) -> None: ... class AsyncSearchCommands: def __init__(self, client: AsyncCommands): ... @@ -2072,22 +2080,18 @@ class AsyncSearchCommands: self, *, name: str, - schema: Union[NestedIndexSchema, FlatIndexSchema], - dataType: str, - prefix: Union[str, List[str]], - language: Optional[str] = None, - skipInitialScan: bool = False, - existsOk: bool = False, - ) -> AsyncSearchIndexCommands: ... - def index( - self, - name: str, - schema: Optional[Union[NestedIndexSchema, FlatIndexSchema]] = None + schema: Schema, + data_type: Union[DataType, str], + prefixes: Union[str, List[str]], + language: Optional[Union[Language, str]] = None, + skip_initial_scan: bool = False, + exists_ok: bool = False, ) -> AsyncSearchIndexCommands: ... + def index(self, name: str) -> AsyncSearchIndexCommands: ... class AsyncSearchIndexCommands: - def __init__(self, client: AsyncCommands, name: str, schema: Optional[Union[NestedIndexSchema, FlatIndexSchema]] = None): ... - async def wait_indexing(self) -> Any: ... + def __init__(self, client: AsyncCommands, name: str): ... + async def wait_indexing(self) -> None: ... async def describe(self) -> IndexDescription: ... async def query( self, @@ -2095,9 +2099,9 @@ class AsyncSearchIndexCommands: filter: Dict[str, Any], limit: Optional[int] = None, offset: Optional[int] = None, - orderBy: Optional[Dict[str, str]] = None, + order_by: Optional[Dict[str, Union[Order, str]]] = None, select: Optional[Dict[str, bool]] = None, - highlight: Optional[Dict[str, Any]] = None, + highlight: Optional[HighlightOptions] = None, ) -> List[QueryResult]: ... - async def count(self, filter: Dict[str, Any]) -> Dict[Literal["count"], int]: ... - async def drop(self) -> int: ... + async def count(self, *, filter: Dict[str, Any]) -> CountResult: ... + async def drop(self) -> None: ... diff --git a/upstash_redis/format.py b/upstash_redis/format.py index 831976e..ad1e04c 100644 --- a/upstash_redis/format.py +++ b/upstash_redis/format.py @@ -1,13 +1,14 @@ from json import loads from typing import Any, Callable, Dict, List, Optional, Tuple, Union -from upstash_redis.typing import RESTResultT -from upstash_redis.utils import GeoSearchResult from upstash_redis.commands import SearchIndexCommands -from upstash_redis.search_utils import ( +from upstash_redis.search import ( deserialize_describe_response, deserialize_query_response, + CountResult, ) +from upstash_redis.typing import RESTResultT +from upstash_redis.utils import GeoSearchResult def to_dict(res: List, _, __) -> Dict: @@ -231,30 +232,29 @@ def format_xpending_response(res, command, _): return res -def format_search_query_response(res: List[Any], _, __) -> List[Dict[str, Any]]: +def format_search_query_response(res: List[Any], _, __): """Format SEARCH.QUERY response into structured results.""" return deserialize_query_response(res) -def format_search_describe_response(res: List[Any], _, __) -> Dict[str, Any]: +def format_search_describe_response(res: List[Any], _, __): """Format SEARCH.DESCRIBE response into index description.""" return deserialize_describe_response(res) -def format_search_count_response(res: Any, _, __) -> Dict[str, int]: +def format_search_count_response(res: Any, _, __): """Format SEARCH.COUNT response.""" - count = int(res) if not isinstance(res, int) else res - return {"count": count} + return CountResult(count=res) def format_search_create_response(res: Any, command: List[str], client: Any): """Format SEARCH.CREATE response by returning SearchIndexCommands instance.""" # Extract index name from command (second parameter) index_name = command[1] if len(command) > 1 else None - + if index_name is None: return res - + return SearchIndexCommands(client, index_name) diff --git a/upstash_redis/search.py b/upstash_redis/search.py new file mode 100644 index 0000000..04a982f --- /dev/null +++ b/upstash_redis/search.py @@ -0,0 +1,220 @@ +"""Type definitions for search index functionality.""" + +import dataclasses +import enum +import json +from typing import Any, Dict, List, Optional, Tuple, TypedDict, Union + + +# Data types that can be indexed +class DataType(str, enum.Enum): + JSON = "JSON" + HASH = "HASH" + STRING = "STRING" + + +class FieldType(str, enum.Enum): + TEXT = "TEXT" + U64 = "U64" + I64 = "I64" + F64 = "F64" + BOOL = "BOOL" + DATE = "DATE" + + +class FieldOptions(TypedDict, total=False): + """Detailed field configuration with options.""" + + type: Union[FieldType, str] + fast: bool + no_tokenize: bool + no_stem: bool + alias: str + + +Schema = Dict[str, Union[Union[FieldType, str], FieldOptions]] + + +class Language(str, enum.Enum): + ENGLISH = "english" + ARABIC = "arabic" + DANISH = "danish" + DUTCH = "dutch" + FINNISH = "finnish" + FRENCH = "french" + GERMAN = "german" + GREEK = "greek" + HUNGARIAN = "hungarian" + ITALIAN = "italian" + NORWEGIAN = "norwegian" + PORTUGUESE = "portuguese" + ROMANIAN = "romanian" + RUSSIAN = "russian" + SPANISH = "spanish" + SWEDISH = "swedish" + TAMIL = "tamil" + TURKISH = "turkish" + + +class Order(str, enum.Enum): + ASC = "ASC" + DESC = "DESC" + + +# Query options +class HighlightOptions(TypedDict, total=False): + """Highlighting options for query results.""" + + fields: List[str] + tags: Tuple[str, str] # "open" and "close" tags + + +# Query result +@dataclasses.dataclass +class QueryResult: + """Result from a query operation.""" + + key: str + score: float + data: Optional[Dict[str, Any]] = None + + +@dataclasses.dataclass +class CountResult: + count: int + + +# Index description +@dataclasses.dataclass +class FieldInfo: + """Field information from describe command.""" + + type: FieldType + alias: Optional[str] + no_stem: bool + no_tokenize: bool + fast: bool + + +@dataclasses.dataclass +class IndexDescription: + """Description of an index.""" + + name: str + data_type: DataType + prefixes: List[str] + language: Language + schema: Dict[str, FieldInfo] + + +def deserialize_query_response(raw_response: List[Any]) -> List[QueryResult]: + """ + Deserialize raw query response into structured results. + + Args: + raw_response: Raw response from SEARCH.QUERY + + Returns: + List of query results + """ + results: List[QueryResult] = [] + + for item in raw_response: + key = item[0] + score = float(item[1]) + data: Optional[Dict[str, Any]] = None + + if len(item) >= 3: + data = {} + for path, value in item[2]: + path_parts = path.split(".") + if len(path_parts) == 1: + data[path] = value + else: + parent = data + for i, part in enumerate(path_parts): + if i < len(path_parts) - 1: + if part not in parent: + parent[part] = {} + + parent = parent[part] + else: + parent[part] = value + + if "$" in data: + data = json.loads(data["$"]) + + result = QueryResult(key=key, score=score, data=data) + results.append(result) + + return results + + +def deserialize_describe_response(raw_response: List[Any]) -> IndexDescription: + """ + Deserialize raw describe response into index description. + + Args: + raw_response: Raw response from SEARCH.DESCRIBE + + Returns: + Index description + """ + + name = "" + data_type: DataType = DataType.JSON + prefixes = [] + language = Language.ENGLISH + schema: Dict[str, FieldInfo] = {} + + for i in range(0, len(raw_response), 2): + descriptor = raw_response[i] + if descriptor == "name": + name = raw_response[i + 1] + elif descriptor == "type": + data_type = DataType[raw_response[i + 1].upper()] + elif descriptor == "prefixes": + prefixes = raw_response[i + 1] + elif descriptor == "language": + language = Language[raw_response[i + 1].upper()] + elif descriptor == "schema": + fields = raw_response[i + 1] + for field_info in fields: + field_name = field_info[0] + field_type = field_info[1] + + alias: Optional[str] = None + no_stem = False + no_tokenize = False + fast = False + + j = 2 + while j < len(field_info): + field_descriptor = field_info[j] + if field_descriptor == "FROM": + alias = field_info[j + 1] + j += 1 + elif field_descriptor == "NOSTEM": + no_stem = True + elif field_descriptor == "NOTOKENIZE": + no_tokenize = True + elif field_descriptor == "FAST": + fast = True + + j += 1 + + schema[field_name] = FieldInfo( + type=field_type, + alias=alias, + no_stem=no_stem, + no_tokenize=no_tokenize, + fast=fast, + ) + + return IndexDescription( + name=name, + data_type=data_type, + prefixes=prefixes, + language=language, + schema=schema, + ) diff --git a/upstash_redis/search_types.py b/upstash_redis/search_types.py deleted file mode 100644 index 4aaae2a..0000000 --- a/upstash_redis/search_types.py +++ /dev/null @@ -1,115 +0,0 @@ -"""Type definitions for search index functionality.""" - -from typing import Any, Dict, List, Literal, Optional, TypedDict, Union - -# Field types -FieldType = Literal["TEXT", "U64", "I64", "F64", "BOOL", "DATE"] - -# Language options for text indexing -Language = Literal[ - "arabic", - "basque", - "catalan", - "danish", - "dutch", - "english", - "finnish", - "french", - "german", - "greek", - "hungarian", - "indonesian", - "irish", - "italian", - "lithuanian", - "nepali", - "norwegian", - "portuguese", - "romanian", - "russian", - "spanish", - "swedish", - "tamil", - "turkish", -] - -# Data types that can be indexed -DataType = Literal["string", "json", "hash"] - - -# Schema definitions -class DetailedField(TypedDict, total=False): - """Detailed field configuration with options.""" - - type: FieldType - fast: bool - noTokenize: bool - noStem: bool - from_: str # Using from_ to avoid Python keyword - - -# Schema can be a nested structure of fields -NestedIndexSchema = Dict[str, Union[FieldType, DetailedField, "NestedIndexSchema"]] -FlatIndexSchema = Dict[str, Union[FieldType, DetailedField]] - -RootQueryFilter = Dict[str, Any] - -# Query options -class HighlightOptions(TypedDict, total=False): - """Highlighting options for query results.""" - - fields: List[str] - tags: Dict[str, str] # "open" and "close" tags - - -class QueryOptions(TypedDict, total=False): - """Options for querying an index.""" - - filter: Any - limit: int - offset: int - orderBy: Dict[str, Literal["ASC", "DESC"]] - select: Dict[str, bool] # Field name to include/exclude - highlight: HighlightOptions - - -# Query result -class QueryResult(TypedDict, total=False): - """Result from a query operation.""" - - key: str - score: str - data: Dict[str, Any] - - -# Index description -class DescribeFieldInfo(TypedDict, total=False): - """Field information from describe command.""" - - type: FieldType - fast: bool - noTokenize: bool - noStem: bool - - -class IndexDescription(TypedDict, total=False): - """Description of an index.""" - - name: str - dataType: DataType - prefixes: List[str] - language: Language - schema: Dict[str, DescribeFieldInfo] - - -# Create index parameters -class CreateIndexParams(TypedDict, total=False): - """Parameters for creating a new index.""" - - name: str - prefix: Union[str, List[str]] - dataType: DataType - schema: Union[NestedIndexSchema, FlatIndexSchema] - language: Language - skipInitialScan: bool - existsOk: bool diff --git a/upstash_redis/search_utils.py b/upstash_redis/search_utils.py deleted file mode 100644 index 61774f2..0000000 --- a/upstash_redis/search_utils.py +++ /dev/null @@ -1,363 +0,0 @@ -"""Utility functions for search index operations.""" - -from typing import Any, Dict, List, Union - -from upstash_redis.search_types import ( - CreateIndexParams, - DescribeFieldInfo, - FieldType, - FlatIndexSchema, - IndexDescription, - NestedIndexSchema, - QueryOptions, - QueryResult, -) - - -class FlattenedField: - """Represents a flattened field from a nested schema.""" - - def __init__( - self, - path: str, - field_type: FieldType, - fast: bool = False, - no_tokenize: bool = False, - no_stem: bool = False, - from_: str = None, - ): - self.path = path - self.field_type = field_type - self.fast = fast - self.no_tokenize = no_tokenize - self.no_stem = no_stem - self.from_ = from_ - - -def _is_field_type(value: Any) -> bool: - """Check if value is a valid field type.""" - return isinstance(value, str) and value in ["TEXT", "U64", "I64", "F64", "BOOL", "DATE"] - - -def _is_detailed_field(value: Any) -> bool: - """Check if value is a detailed field configuration.""" - return isinstance(value, dict) and "type" in value and _is_field_type(value.get("type")) - - -def _is_nested_schema(value: Any) -> bool: - """Check if value is a nested schema.""" - return isinstance(value, dict) and not _is_detailed_field(value) and not _is_field_type(value) - - -def flatten_schema( - schema: Union[NestedIndexSchema, FlatIndexSchema], path_prefix: List[str] = None -) -> List[FlattenedField]: - """ - Flatten a nested schema into a list of field definitions. - - Args: - schema: The schema to flatten - path_prefix: Current path prefix for nested fields - - Returns: - List of flattened field definitions - """ - if path_prefix is None: - path_prefix = [] - - fields: List[FlattenedField] = [] - - for key, value in schema.items(): - current_path = path_prefix + [key] - path_string = ".".join(current_path) - - if _is_field_type(value): - fields.append(FlattenedField(path=path_string, field_type=value)) - elif _is_detailed_field(value): - detailed = value - fields.append( - FlattenedField( - path=path_string, - field_type=detailed["type"], - fast=detailed.get("fast", False), - no_tokenize=detailed.get("noTokenize", False), - no_stem=detailed.get("noStem", False), - from_=detailed.get("from_"), - ) - ) - elif _is_nested_schema(value): - nested_fields = flatten_schema(value, current_path) - fields.extend(nested_fields) - - return fields - - -def build_create_index_command(params: CreateIndexParams) -> List[str]: - """ - Build the SEARCH.CREATE command from parameters. - - Args: - params: Create index parameters - - Returns: - Command list ready to execute - """ - command: List[str] = ["SEARCH.CREATE", params["name"]] - - # Add options BEFORE "ON" keyword - if params.get("skipInitialScan", False): - command.append("SKIPINITIALSCAN") - - if params.get("existsOk", False): - command.append("EXISTSOK") - - # Add "ON" keyword and data type - command.append("ON") - data_type = params["dataType"].upper() - command.append(data_type) - - # Add prefixes - prefixes = params["prefix"] - if isinstance(prefixes, str): - prefixes = [prefixes] - - command.append("PREFIX") - command.append(str(len(prefixes))) - command.extend(prefixes) - - # Add language if specified - if "language" in params: - command.extend(["LANGUAGE", params["language"]]) - - # Add schema - schema = params["schema"] - flattened = flatten_schema(schema) - - command.append("SCHEMA") - for field in flattened: - # Add field path or from source - if field.from_: - command.extend(["AS", field.path, "FROM", field.from_]) - else: - command.append(field.path) - - # Add field type - command.append(field.field_type) - - # Add field options - if field.no_tokenize: - command.append("NOTOKENIZE") - if field.no_stem: - command.append("NOSTEM") - if field.fast: - command.append("FAST") - - return command - - -def build_query_command( - command_name: str, index_name: str, options: QueryOptions = None -) -> List[str]: - """ - Build SEARCH.QUERY or SEARCH.COUNT command. - - Args: - command_name: "SEARCH.QUERY" or "SEARCH.COUNT" - index_name: Name of the index - options: Query options - - Returns: - Command list ready to execute - """ - import json - - command: List[str] = [command_name, index_name] - - # Serialize filter to JSON (compact format without spaces) - filter_json = "{}" - if options and "filter" in options: - filter_json = json.dumps(options["filter"], separators=(',', ':')) - - command.append(filter_json) - - if not options: - return command - - # Add limit - if "limit" in options: - command.extend(["LIMIT", str(options["limit"])]) - - # Add offset - if "offset" in options: - command.extend(["OFFSET", str(options["offset"])]) - - # Add orderBy - if "orderBy" in options: - for field, direction in options["orderBy"].items(): - command.extend(["SORTBY", field, direction]) - - # Add highlight - if "highlight" in options: - highlight = options["highlight"] - command.append("HIGHLIGHT") - if "fields" in highlight: - fields = highlight["fields"] - command.extend(["FIELDS", str(len(fields))]) - command.extend(fields) - if "preTag" in highlight and "postTag" in highlight: - command.extend(["TAGS", highlight["preTag"], highlight["postTag"]]) - - # Add select - if "select" in options: - select = options["select"] - if not select: # Empty dict means NOCONTENT - command.append("NOCONTENT") - else: - command.append("RETURN") - fields = [f for f, include in select.items() if include] - command.append(str(len(fields))) - command.extend(fields) - - return command - - # Add highlight - if "highlight" in options: - highlight = options["highlight"] - command.append("HIGHLIGHT") - - if "fields" in highlight and highlight["fields"]: - command.append("FIELDS") - command.append(str(len(highlight["fields"]))) - command.extend(highlight["fields"]) - - if "tags" in highlight: - tags = highlight["tags"] - if "open" in tags and "close" in tags: - command.extend(["TAGS", tags["open"], tags["close"]]) - - return command - - -def deserialize_query_response(raw_response: List[Any]) -> List[QueryResult]: - """ - Deserialize raw query response into structured results. - - Args: - raw_response: Raw response from SEARCH.QUERY - - Returns: - List of query results - """ - results: List[QueryResult] = [] - - for item_raw in raw_response: - if not isinstance(item_raw, (list, tuple)) or len(item_raw) < 2: - continue - - key = item_raw[0] - score = item_raw[1] - - result: QueryResult = {"key": key, "score": score} - - if len(item_raw) > 2: - raw_fields = item_raw[2] - - if isinstance(raw_fields, (list, tuple)) and len(raw_fields) > 0: - data: Dict[str, Any] = {} - - # Process field pairs - for field_raw in raw_fields: - if not isinstance(field_raw, (list, tuple)) or len(field_raw) < 2: - continue - - field_key = field_raw[0] - field_value = field_raw[1] - - # Handle nested paths - path_parts = field_key.split(".") - if len(path_parts) == 1: - data[field_key] = field_value - else: - # Build nested structure - current_obj = data - for i, part in enumerate(path_parts[:-1]): - if part not in current_obj: - current_obj[part] = {} - current_obj = current_obj[part] - current_obj[path_parts[-1]] = field_value - - # If $ key exists (full document), use its contents - if "$" in data: - dollar_value = data["$"] - # Parse JSON string if needed - if isinstance(dollar_value, str): - import json - try: - data = json.loads(dollar_value) - except (json.JSONDecodeError, ValueError): - # If parsing fails, use the string as-is - data = dollar_value - else: - data = dollar_value - - result["data"] = data - - results.append(result) - - return results - - -def deserialize_describe_response(raw_response: List[Any]) -> IndexDescription: - """ - Deserialize raw describe response into index description. - - Args: - raw_response: Raw response from SEARCH.DESCRIBE - - Returns: - Index description - """ - description: IndexDescription = {} - - i = 0 - while i < len(raw_response): - descriptor = raw_response[i] - value = raw_response[i + 1] if i + 1 < len(raw_response) else None - - if descriptor == "name": - description["name"] = value - elif descriptor == "type": - if value: - description["dataType"] = value.lower() - elif descriptor == "prefixes": - description["prefixes"] = value if isinstance(value, list) else [] - elif descriptor == "language": - description["language"] = value - elif descriptor == "schema": - schema: Dict[str, DescribeFieldInfo] = {} - if isinstance(value, list): - for field_desc in value: - if not isinstance(field_desc, list) or len(field_desc) < 2: - continue - - field_name = field_desc[0] - field_info: DescribeFieldInfo = {"type": field_desc[1]} - - # Parse field options - for j in range(2, len(field_desc)): - option = field_desc[j] - if option == "NOSTEM": - field_info["noStem"] = True - elif option == "NOTOKENIZE": - field_info["noTokenize"] = True - elif option == "FAST": - field_info["fast"] = True - - schema[field_name] = field_info - - description["schema"] = schema - - i += 2 - - return description - From ea1dcb60d8561765493bb4ca52be085bc9470e0b Mon Sep 17 00:00:00 2001 From: CahidArda Date: Fri, 30 Jan 2026 16:36:30 +0300 Subject: [PATCH 09/12] fix: add --- examples/search.py | 11 +++++++++-- .../asyncio/search/test_search_index_async.py | 8 ++++++++ tests/commands/search/test_search_advanced.py | 7 +++++++ tests/commands/search/test_search_index.py | 11 +++++++++++ 4 files changed, 35 insertions(+), 2 deletions(-) diff --git a/examples/search.py b/examples/search.py index bf8614b..e7e9c99 100644 --- a/examples/search.py +++ b/examples/search.py @@ -292,14 +292,21 @@ for result in results: print(f" - {result.data['name']}") +# Smart search (intelligent matching) +print("\n2. Smart search for 'lapto' (partial term):") +results = products_index.query(filter={"name": {"$smart": "lapto"}}) +print(f" Found {len(results)} results with smart matching") +for result in results: + print(f" - {result.data['name']}") + # Phrase search -print("\n2. Phrase search for 'Wireless Mouse':") +print("\n3. Phrase search for 'Wireless Mouse':") results = products_index.query(filter={"name": {"$phrase": "Wireless Mouse"}}) for result in results: print(f" Found: {result.data['name']}") # Range query -print("\n3. Products priced between $10 and $100:") +print("\n4. Products priced between $10 and $100:") results = products_index.query( filter={"price": {"$gte": 10, "$lte": 100}}, order_by={"price": "ASC"}, diff --git a/tests/commands/asyncio/search/test_search_index_async.py b/tests/commands/asyncio/search/test_search_index_async.py index b9a644a..4ecfc70 100644 --- a/tests/commands/asyncio/search/test_search_index_async.py +++ b/tests/commands/asyncio/search/test_search_index_async.py @@ -144,6 +144,14 @@ async def test_query_text_fuzzy(self, query_index: QueryIndexFixture): # Fuzzy search should find Laptop variants despite typo assert len(result) >= 2 + async def test_query_text_smart(self, query_index: QueryIndexFixture): + """Test async querying with smart search.""" + index = query_index.index + result = await index.query(filter={"name": {"$smart": "lapto"}}) + + # Smart search should find Laptop variants with partial term + assert len(result) >= 2 + async def test_query_text_phrase(self, query_index: QueryIndexFixture): """Test async querying with phrase matching.""" index = query_index.index diff --git a/tests/commands/search/test_search_advanced.py b/tests/commands/search/test_search_advanced.py index 3eadfa1..9dc9611 100644 --- a/tests/commands/search/test_search_advanced.py +++ b/tests/commands/search/test_search_advanced.py @@ -124,6 +124,13 @@ def test_query_json_fuzzy(self, json_index: IndexFixture): assert len(result) > 0 + def test_query_json_smart(self, json_index: IndexFixture): + """Test querying JSON index with smart search.""" + index = json_index.index + result = index.query(filter={"name": {"$smart": "lapto"}}) + + assert len(result) > 0 + def test_query_json_phrase(self, json_index: IndexFixture): """Test querying JSON index with phrase matching.""" index = json_index.index diff --git a/tests/commands/search/test_search_index.py b/tests/commands/search/test_search_index.py index 955d8f6..ec27ef4 100644 --- a/tests/commands/search/test_search_index.py +++ b/tests/commands/search/test_search_index.py @@ -287,6 +287,17 @@ def test_query_text_fuzzy(self, string_index: StringIndexFixture): assert r.key.startswith(string_index.name) assert r.score > 0.0 + def test_query_text_smart(self, string_index: StringIndexFixture): + """Test querying with smart search for intelligent matching.""" + index = string_index.index + result = index.query(filter={"name": {"$smart": "lapto"}}) + + # Smart search should find Laptop variants with partial term + assert len(result) >= 2 + for r in result: + assert r.key.startswith(string_index.name) + assert r.score > 0.0 + def test_query_text_phrase(self, string_index: StringIndexFixture): """Test querying with phrase matching.""" index = string_index.index From 03bf0fac30a74d29b63a92a0823e2b3ea0ce83d2 Mon Sep 17 00:00:00 2001 From: CahidArda Date: Fri, 30 Jan 2026 18:21:14 +0300 Subject: [PATCH 10/12] feat: add score_func param --- .../search/test_search_scorefunc_async.py | 165 ++++++++++++++ .../commands/search/test_search_scorefunc.py | 214 ++++++++++++++++++ upstash_redis/commands.py | 7 + upstash_redis/commands.pyi | 3 + upstash_redis/search.py | 59 +++++ upstash_redis/utils.py | 40 +++- 6 files changed, 487 insertions(+), 1 deletion(-) create mode 100644 tests/commands/asyncio/search/test_search_scorefunc_async.py create mode 100644 tests/commands/search/test_search_scorefunc.py diff --git a/tests/commands/asyncio/search/test_search_scorefunc_async.py b/tests/commands/asyncio/search/test_search_scorefunc_async.py new file mode 100644 index 0000000..56e995c --- /dev/null +++ b/tests/commands/asyncio/search/test_search_scorefunc_async.py @@ -0,0 +1,165 @@ +"""Tests for async score function queries.""" + +import pytest +import pytest_asyncio + +from upstash_redis.asyncio import Redis as AsyncRedis + + +def random_id() -> str: + """Generate a random ID for test isolation.""" + import random + import string + + return "".join(random.choices(string.ascii_lowercase + string.digits, k=8)) + + +@pytest.fixture +def async_redis_client() -> AsyncRedis: + """Create an async Redis client for testing.""" + return AsyncRedis.from_env() + + +@pytest.fixture +def cleanup_indexes(): + """Track and cleanup indexes created during tests.""" + indexes = [] + + yield indexes + + # Cleanup after test + import asyncio + + async def cleanup(): + redis = AsyncRedis.from_env() + for index_name in indexes: + try: + index = redis.search.index(index_name) + await index.drop() + except Exception: + pass + + asyncio.run(cleanup()) + + +@pytest.mark.asyncio +class TestAsyncScoreFunctionQueries: + """Tests for async querying with score functions.""" + + @pytest_asyncio.fixture(scope="class") + async def scorefunc_index(self) -> dict: + """Create a test index with data for score function querying.""" + redis = AsyncRedis.from_env() + name = f"test-scorefunc-async-{random_id()}" + prefix = f"{name}:" + keys = [] + + try: + # Create index + index = await redis.search.create_index( + name=name, + schema={ + "name": "TEXT", + "popularity": {"type": "U64", "fast": True}, + "recency": {"type": "U64", "fast": True}, + }, + data_type="hash", + prefixes=prefix, + ) + + # Add test data with different popularity scores + products = [ + {"name": "Laptop Pro", "popularity": "1000", "recency": "100"}, + {"name": "Laptop Basic", "popularity": "500", "recency": "200"}, + {"name": "Laptop Air", "popularity": "2000", "recency": "50"}, + ] + + for i, product in enumerate(products): + key = f"{prefix}{i}" + keys.append(key) + await redis.hset(key, values=product) + + # Wait for indexing + await index.wait_indexing() + + yield {"index": index, "redis": redis, "keys": keys, "name": name} + + finally: + # Cleanup + try: + await index.drop() + except Exception: + pass + if keys: + await redis.delete(*keys) + + async def test_async_query_with_simple_scorefunc(self, scorefunc_index: dict): + """Test async querying with simple scoreFunc that boosts by popularity.""" + index = scorefunc_index["index"] + result = await index.query( + filter={"name": {"$eq": "Laptop"}}, + score_func="popularity", + ) + + assert len(result) == 3 + # Higher popularity should result in higher scores + assert result[0].data is not None + assert result[0].data["name"] == "Laptop Air" # popularity 2000 + + async def test_async_query_with_scorefunc_using_modifier( + self, scorefunc_index: dict + ): + """Test async querying with scoreFunc using modifier.""" + index = scorefunc_index["index"] + result = await index.query( + filter={"name": {"$eq": "Laptop"}}, + score_func={ + "field": "popularity", + "modifier": "LOG1P", + "factor": 2, + }, + ) + + assert len(result) == 3 + # Results should be affected by log1p(popularity) * 2.0 + assert result[0].data is not None + assert result[0].data["name"] == "Laptop Air" + + async def test_async_query_with_scorefunc_using_scoremode_replace( + self, scorefunc_index: dict + ): + """Test async querying with scoreFunc using scoreMode replace.""" + index = scorefunc_index["index"] + result = await index.query( + filter={"name": {"$eq": "Laptop"}}, + score_func={ + "field": "popularity", + "scoreMode": "REPLACE", + }, + ) + + assert len(result) == 3 + # Scores should be replaced with popularity values + assert result[0].data is not None + assert result[0].data["name"] == "Laptop Air" # popularity 2000 + assert result[0].score > result[1].score + + async def test_async_query_with_multiple_field_values_combined( + self, scorefunc_index: dict + ): + """Test async querying with multiple field values combined.""" + index = scorefunc_index["index"] + result = await index.query( + filter={"name": {"$eq": "Laptop"}}, + score_func={ + "fields": [ + {"field": "popularity", "modifier": "LOG1P"}, + {"field": "recency", "modifier": "LOG1P"}, + ], + "combineMode": "SUM", + }, + ) + + assert len(result) == 3 + # Results should be affected by log1p(popularity) + log1p(recency) + assert result[0].data is not None diff --git a/tests/commands/search/test_search_scorefunc.py b/tests/commands/search/test_search_scorefunc.py new file mode 100644 index 0000000..16eaf75 --- /dev/null +++ b/tests/commands/search/test_search_scorefunc.py @@ -0,0 +1,214 @@ +"""Tests for score function queries.""" + +import math +from typing import Generator, List, TypedDict + +import pytest + +from upstash_redis import Redis +from upstash_redis.commands import SearchIndexCommands + + +class ScoreFuncFixture(TypedDict): + index: SearchIndexCommands + redis: Redis + keys: List[str] + name: str + + +def random_id() -> str: + """Generate a random ID for test isolation.""" + import random + import string + + return "".join(random.choices(string.ascii_lowercase + string.digits, k=8)) + + +@pytest.fixture +def redis_client() -> Redis: + """Create a Redis client for testing.""" + return Redis.from_env() + + +@pytest.fixture +def cleanup_indexes(): + """Track and cleanup indexes created during tests.""" + indexes = [] + + yield indexes + + # Cleanup after test + redis = Redis.from_env() + for index_name in indexes: + try: + index = redis.search.index(index_name) + index.drop() + except Exception: + pass + + +class TestScoreFunctionQueries: + """Tests for querying with score functions.""" + + @pytest.fixture(scope="class") + def scorefunc_index(self) -> Generator[ScoreFuncFixture, None, None]: + """Create a test index with data for score function querying.""" + redis = Redis.from_env() + name = f"test-scorefunc-{random_id()}" + prefix = f"{name}:" + keys = [] + + try: + # Create index + index = redis.search.create_index( + name=name, + schema={ + "name": "TEXT", + "popularity": {"type": "U64", "fast": True}, + "recency": {"type": "U64", "fast": True}, + }, + data_type="hash", + prefixes=prefix, + ) + + # Add test data with different popularity scores + products = [ + {"name": "Laptop Pro", "popularity": "1000", "recency": "100"}, + {"name": "Laptop Basic", "popularity": "500", "recency": "200"}, + {"name": "Laptop Air", "popularity": "2000", "recency": "50"}, + {"name": "Laptop Missing", "recency": "150"}, + ] + + for i, product in enumerate(products): + key = f"{prefix}{i}" + keys.append(key) + redis.hset(key, values=product) + + # Wait for indexing + index.wait_indexing() + + yield {"index": index, "redis": redis, "keys": keys, "name": name} + + finally: + # Cleanup + try: + index.drop() + except Exception: + pass + if keys: + redis.delete(*keys) + + def test_query_with_simple_scorefunc(self, scorefunc_index: ScoreFuncFixture): + """Test querying with simple scoreFunc that boosts by popularity.""" + index = scorefunc_index["index"] + result = index.query( + filter={"name": {"$eq": "Laptop"}}, + score_func={"field": "popularity", "scoreMode": "REPLACE"}, + ) + + assert len(result) == 4 + # Scores should exactly match the popularity values + scores_by_name = {r.data["name"]: r.score for r in result if r.data} + assert scores_by_name["Laptop Air"] == pytest.approx(2000.0) + assert scores_by_name["Laptop Pro"] == pytest.approx(1000.0) + assert scores_by_name["Laptop Basic"] == pytest.approx(500.0) + + def test_query_with_scorefunc_using_modifier(self, scorefunc_index: ScoreFuncFixture): + """Test querying with scoreFunc using modifier.""" + index = scorefunc_index["index"] + result = index.query( + filter={"name": {"$eq": "Laptop"}}, + score_func={ + "field": "popularity", + "modifier": "LOG1P", + "factor": 2, + "scoreMode": "REPLACE", + }, + ) + + assert len(result) == 4 + scores_by_name = {r.data["name"]: r.score for r in result if r.data} + assert scores_by_name["Laptop Air"] == pytest.approx(2 * math.log10(1 + 2000)) + assert scores_by_name["Laptop Pro"] == pytest.approx(2 * math.log10(1 + 1000)) + assert scores_by_name["Laptop Basic"] == pytest.approx(2 * math.log10(1 + 500)) + + def test_query_with_scorefunc_using_scoremode_replace( + self, scorefunc_index: ScoreFuncFixture + ): + """Test querying with scoreFunc using scoreMode replace.""" + index = scorefunc_index["index"] + result = index.query( + filter={"name": {"$eq": "Laptop"}}, + score_func={ + "field": "popularity", + "scoreMode": "REPLACE", + }, + ) + + assert len(result) == 4 + scores_by_name = {r.data["name"]: r.score for r in result if r.data} + assert scores_by_name["Laptop Air"] == pytest.approx(2000.0) + assert scores_by_name["Laptop Pro"] == pytest.approx(1000.0) + assert scores_by_name["Laptop Basic"] == pytest.approx(500.0) + + def test_query_with_multiple_field_values_combined(self, scorefunc_index: ScoreFuncFixture): + """Test querying with multiple field values combined.""" + index = scorefunc_index["index"] + result = index.query( + filter={"name": {"$eq": "Laptop"}}, + score_func={ + "fields": [ + {"field": "popularity", "modifier": "LOG1P"}, + {"field": "recency", "modifier": "LOG1P"}, + ], + "combineMode": "SUM", + "scoreMode": "REPLACE", + }, + ) + + assert len(result) == 4 + scores_by_name = {r.data["name"]: r.score for r in result if r.data} + assert scores_by_name["Laptop Air"] == pytest.approx( + math.log10(1 + 2000) + math.log10(1 + 50) + ) + assert scores_by_name["Laptop Pro"] == pytest.approx( + math.log10(1 + 1000) + math.log10(1 + 100) + ) + assert scores_by_name["Laptop Basic"] == pytest.approx( + math.log10(1 + 500) + math.log10(1 + 200) + ) + + def test_query_with_scorefunc_modifier_and_missing(self, scorefunc_index: ScoreFuncFixture): + """Test querying with scoreFunc using modifier and missing value.""" + index = scorefunc_index["index"] + result = index.query( + filter={"name": {"$eq": "Laptop"}}, + score_func={ + "field": "popularity", + "modifier": "LOG1P", + "missing": 1, + "scoreMode": "REPLACE", + }, + ) + + assert len(result) == 4 + scores_by_name = {r.data["name"]: r.score for r in result if r.data} + assert scores_by_name["Laptop Missing"] == pytest.approx(math.log10(1 + 1)) + + def test_query_with_multiple_fields_and_scoremode(self, scorefunc_index: ScoreFuncFixture): + """Test querying with multiple field values and scoreMode.""" + index = scorefunc_index["index"] + result = index.query( + filter={"name": {"$eq": "Laptop"}}, + score_func={ + "fields": ["popularity", "recency"], + "combineMode": "MULTIPLY", + "scoreMode": "REPLACE", + }, + ) + + assert len(result) == 4 + scores_by_name = {r.data["name"]: r.score for r in result if r.data} + assert scores_by_name["Laptop Air"] == pytest.approx(2000 * 50) + assert scores_by_name["Laptop Pro"] == pytest.approx(1000 * 100) + assert scores_by_name["Laptop Basic"] == pytest.approx(500 * 200) diff --git a/upstash_redis/commands.py b/upstash_redis/commands.py index a937d64..2313b2a 100644 --- a/upstash_redis/commands.py +++ b/upstash_redis/commands.py @@ -19,9 +19,11 @@ FieldType, Order, HighlightOptions, + ScoreFunc, ) from upstash_redis.typing import FloatMinMaxT, JSONValueT, ValueT from upstash_redis.utils import ( + build_score_func, handle_georadius_write_exceptions, handle_geosearch_exceptions, handle_non_deprecated_zrange_exceptions, @@ -5596,6 +5598,7 @@ def query( order_by: Optional[Dict[str, Union[Order, str]]] = None, select: Optional[Dict[str, bool]] = None, highlight: Optional[HighlightOptions] = None, + score_func: Optional[ScoreFunc] = None, ) -> ResponseT: """ Query the index with filters and options. @@ -5635,6 +5638,10 @@ def query( open_tag, close_tag = highlight["tags"] command.extend(("TAGS", open_tag, close_tag)) + if score_func: + command.append("SCOREFUNC") + build_score_func(command, score_func) + return self.client.execute(command) def count(self, *, filter: Dict[str, Any]) -> ResponseT: diff --git a/upstash_redis/commands.pyi b/upstash_redis/commands.pyi index 879d679..335606f 100644 --- a/upstash_redis/commands.pyi +++ b/upstash_redis/commands.pyi @@ -10,6 +10,7 @@ from upstash_redis.search import ( Order, HighlightOptions, CountResult, + ScoreFunc, ) from upstash_redis.typing import FloatMinMaxT, ValueT, JSONValueT from upstash_redis.utils import GeoSearchResult @@ -2070,6 +2071,7 @@ class SearchIndexCommands: order_by: Optional[Dict[str, Union[Order, str]]] = None, select: Optional[Dict[str, bool]] = None, highlight: Optional[HighlightOptions] = None, + score_func: Optional[ScoreFunc] = None, ) -> List[QueryResult]: ... def count(self, *, filter: Dict[str, Any]) -> CountResult: ... def drop(self) -> None: ... @@ -2102,6 +2104,7 @@ class AsyncSearchIndexCommands: order_by: Optional[Dict[str, Union[Order, str]]] = None, select: Optional[Dict[str, bool]] = None, highlight: Optional[HighlightOptions] = None, + score_func: Optional[ScoreFunc] = None, ) -> List[QueryResult]: ... async def count(self, *, filter: Dict[str, Any]) -> CountResult: ... async def drop(self) -> None: ... diff --git a/upstash_redis/search.py b/upstash_redis/search.py index 04a982f..27390aa 100644 --- a/upstash_redis/search.py +++ b/upstash_redis/search.py @@ -61,6 +61,65 @@ class Order(str, enum.Enum): DESC = "DESC" +class ScoreModifier(str, enum.Enum): + """Score modifiers for field values.""" + + NONE = "NONE" + LOG = "LOG" + LOG1P = "LOG1P" + LOG2P = "LOG2P" + LN = "LN" + LN1P = "LN1P" + LN2P = "LN2P" + SQUARE = "SQUARE" + SQRT = "SQRT" + RECIPROCAL = "RECIPROCAL" + +class ScoreMode(str, enum.Enum): + """Score combination modes.""" + + SUM = "SUM" + MULTIPLY = "MULTIPLY" + REPLACE = "REPLACE" + + +class CombineMode(str, enum.Enum): + """Combine modes for multiple field values.""" + + SUM = "SUM" + MULTIPLY = "MULTIPLY" + + +class ScoreByField(TypedDict, total=False): + """Score function configuration for a single field.""" + + field: str + modifier: ScoreModifier + factor: Union[int, float] + missing: Union[int, float] + scoreMode: ScoreMode + + +class MultipleField(TypedDict, total=False): + """Score function configuration for a field in multiple fields.""" + + field: str + modifier: ScoreModifier + factor: Union[int, float] + missing: Union[int, float] + + +class ScoreByMultipleFields(TypedDict, total=False): + """Score function configuration for multiple fields.""" + + fields: List[Union[str, MultipleField]] + combineMode: CombineMode + scoreMode: ScoreMode + + +ScoreFunc = Union[str, ScoreByField, ScoreByMultipleFields] + + # Query options class HighlightOptions(TypedDict, total=False): """Highlighting options for query results.""" diff --git a/upstash_redis/utils.py b/upstash_redis/utils.py index 0f0628a..aafbbe0 100644 --- a/upstash_redis/utils.py +++ b/upstash_redis/utils.py @@ -1,6 +1,7 @@ from dataclasses import dataclass -from typing import Any, Literal, Optional +from typing import Any, List, Literal, Optional, Dict +from upstash_redis.search import ScoreFunc from upstash_redis.typing import FloatMinMaxT @@ -127,3 +128,40 @@ def handle_zrangebylex_exceptions( if number_are_not_none(offset, count, number=1): raise Exception('Both "offset" and "count" must be specified.') + + +def build_score_func(command: List, score_func: ScoreFunc) -> None: + """Build the SCOREFUNC portion of the command.""" + if isinstance(score_func, str): + # Simple field name + command.extend(("FIELDVALUE", score_func)) + elif "fields" in score_func: + # Multiple fields + if "combineMode" in score_func: + command.extend(("COMBINEMODE", score_func["combineMode"].upper())) + if "scoreMode" in score_func: + command.extend(("SCOREMODE", score_func["scoreMode"].upper())) + + for field_spec in score_func["fields"]: + if isinstance(field_spec, str): + command.extend(("FIELDVALUE", field_spec)) + else: + build_field_value(command, field_spec) + else: + # Single field with options + if "scoreMode" in score_func: + command.extend(("SCOREMODE", score_func["scoreMode"].upper())) + build_field_value(command, score_func) + +def build_field_value( + command: List, field_spec: Dict[str, Any] +) -> None: + """Build a FIELDVALUE portion with modifiers.""" + command.extend(("FIELDVALUE", field_spec["field"])) + + if "modifier" in field_spec: + command.extend(("MODIFIER", field_spec["modifier"].upper())) + if "factor" in field_spec: + command.extend(("FACTOR", field_spec["factor"])) + if "missing" in field_spec: + command.extend(("MISSING", field_spec["missing"])) \ No newline at end of file From 54efe0d95c0d17fdb53e82f2107670601801361f Mon Sep 17 00:00:00 2001 From: CahidArda Date: Wed, 18 Feb 2026 22:42:10 +0300 Subject: [PATCH 11/12] feat: add aggregate, alias, KEYWORD/FACET field types and fix mypy errors Add aggregation framework (SEARCH.AGGREGATE) with response parsing for metric, stats, bucket, and nested sub-aggregations. Add alias commands (SEARCH.ALIASADD, SEARCH.ALIASDEL, SEARCH.LISTALIASES) at both top-level and index-level. Add KEYWORD and FACET field types with $in operator support. Update formatters for new commands and handle not-found cases in query/describe. Fix all pre-existing mypy errors in scorefunc tests and utils by using proper enum types and cast narrowing. Co-Authored-By: Claude Opus 4.6 --- .../search/test_search_aggregate_async.py | 275 +++++++++++++++ .../asyncio/search/test_search_alias_async.py | 136 ++++++++ .../search/test_search_field_types_async.py | 161 +++++++++ .../search/test_search_scorefunc_async.py | 15 +- .../commands/search/test_search_aggregate.py | 317 ++++++++++++++++++ tests/commands/search/test_search_alias.py | 134 ++++++++ .../search/test_search_field_types.py | 158 +++++++++ .../commands/search/test_search_scorefunc.py | 41 ++- upstash_redis/commands.py | 74 ++++ upstash_redis/commands.pyi | 42 ++- upstash_redis/format.py | 42 ++- upstash_redis/search.py | 106 ++++++ upstash_redis/utils.py | 31 +- 13 files changed, 1487 insertions(+), 45 deletions(-) create mode 100644 tests/commands/asyncio/search/test_search_aggregate_async.py create mode 100644 tests/commands/asyncio/search/test_search_alias_async.py create mode 100644 tests/commands/asyncio/search/test_search_field_types_async.py create mode 100644 tests/commands/search/test_search_aggregate.py create mode 100644 tests/commands/search/test_search_alias.py create mode 100644 tests/commands/search/test_search_field_types.py diff --git a/tests/commands/asyncio/search/test_search_aggregate_async.py b/tests/commands/asyncio/search/test_search_aggregate_async.py new file mode 100644 index 0000000..2b39e5f --- /dev/null +++ b/tests/commands/asyncio/search/test_search_aggregate_async.py @@ -0,0 +1,275 @@ +"""Async tests for search aggregation functionality.""" + +import json +import random +import string + +import pytest +import pytest_asyncio + +from upstash_redis.asyncio import Redis + + +def random_id() -> str: + return "".join(random.choices(string.ascii_lowercase + string.digits, k=8)) + + +@pytest_asyncio.fixture +async def async_redis_client(): + return Redis.from_env() + + +@pytest.mark.asyncio +class TestAsyncAggregation: + """Async tests for aggregation queries.""" + + @pytest_asyncio.fixture + async def agg_index(self, async_redis_client): + name = f"test-async-agg-{random_id()}" + prefix = f"{name}:" + keys = [] + + redis = async_redis_client + + schema = { + "name": "TEXT", + "category": "KEYWORD", + "price": {"type": "F64", "fast": True}, + "stock": {"type": "U64", "fast": True}, + "active": "BOOL", + } + + index = await redis.search.create_index( + name=name, schema=schema, data_type="string", prefixes=prefix + ) + + test_data = [ + { + "name": "Laptop Pro", + "category": "electronics", + "price": 1299.99, + "stock": 50, + "active": True, + }, + { + "name": "Laptop Basic", + "category": "electronics", + "price": 599.99, + "stock": 100, + "active": True, + }, + { + "name": "Wireless Mouse", + "category": "electronics", + "price": 29.99, + "stock": 200, + "active": True, + }, + { + "name": "USB Cable", + "category": "accessories", + "price": 9.99, + "stock": 500, + "active": True, + }, + { + "name": "Phone Case", + "category": "accessories", + "price": 19.99, + "stock": 300, + "active": False, + }, + { + "name": "Keyboard", + "category": "electronics", + "price": 79.99, + "stock": 150, + "active": True, + }, + { + "name": "Monitor", + "category": "electronics", + "price": 399.99, + "stock": 75, + "active": True, + }, + { + "name": "Headphones", + "category": "audio", + "price": 149.99, + "stock": 120, + "active": True, + }, + { + "name": "Speaker", + "category": "audio", + "price": 59.99, + "stock": 80, + "active": True, + }, + ] + + for i, datum in enumerate(test_data): + key = f"{prefix}{i}" + keys.append(key) + await redis.set(key, json.dumps(datum)) + + await index.wait_indexing() + + yield {"index": index, "redis": redis, "keys": keys, "name": name} + + try: + await index.drop() + except Exception: + pass + if keys: + await redis.delete(*keys) + + async def test_avg_aggregation(self, agg_index): + index = agg_index["index"] + result = await index.aggregate( + aggregations={"avg_price": {"$avg": {"field": "price"}}} + ) + + assert "avg_price" in result + assert "value" in result["avg_price"] + + async def test_sum_aggregation(self, agg_index): + index = agg_index["index"] + result = await index.aggregate( + aggregations={"total_price": {"$sum": {"field": "price"}}} + ) + + assert "total_price" in result + assert "value" in result["total_price"] + + async def test_min_aggregation(self, agg_index): + index = agg_index["index"] + result = await index.aggregate( + aggregations={"min_price": {"$min": {"field": "price"}}} + ) + + assert "min_price" in result + assert "value" in result["min_price"] + + async def test_max_aggregation(self, agg_index): + index = agg_index["index"] + result = await index.aggregate( + aggregations={"max_price": {"$max": {"field": "price"}}} + ) + + assert "max_price" in result + assert "value" in result["max_price"] + + async def test_count_aggregation(self, agg_index): + index = agg_index["index"] + result = await index.aggregate( + aggregations={"doc_count": {"$count": {"field": "price"}}} + ) + + assert "doc_count" in result + assert "value" in result["doc_count"] + + async def test_stats_aggregation(self, agg_index): + index = agg_index["index"] + result = await index.aggregate( + aggregations={"price_stats": {"$stats": {"field": "price"}}} + ) + + assert "price_stats" in result + stats = result["price_stats"] + assert "count" in stats + assert "min" in stats + assert "max" in stats + + async def test_terms_aggregation(self, agg_index): + index = agg_index["index"] + result = await index.aggregate( + aggregations={"by_category": {"$terms": {"field": "category"}}} + ) + + assert "by_category" in result + assert "buckets" in result["by_category"] + buckets = result["by_category"]["buckets"] + assert len(buckets) > 0 + + async def test_range_aggregation(self, agg_index): + index = agg_index["index"] + result = await index.aggregate( + aggregations={ + "price_ranges": { + "$range": { + "field": "price", + "ranges": [ + {"to": 50}, + {"from": 50, "to": 500}, + {"from": 500}, + ], + } + } + } + ) + + assert "price_ranges" in result + assert "buckets" in result["price_ranges"] + + async def test_histogram_aggregation(self, agg_index): + index = agg_index["index"] + result = await index.aggregate( + aggregations={ + "price_hist": { + "$histogram": { + "field": "price", + "interval": 200, + } + } + } + ) + + assert "price_hist" in result + assert "buckets" in result["price_hist"] + + async def test_nested_sub_aggregation(self, agg_index): + index = agg_index["index"] + result = await index.aggregate( + aggregations={ + "by_category": { + "$terms": { + "field": "category", + }, + "$aggs": { + "avg_price": {"$avg": {"field": "price"}}, + }, + } + } + ) + + assert "by_category" in result + assert "buckets" in result["by_category"] + for bucket in result["by_category"]["buckets"]: + assert "key" in bucket + assert "avg_price" in bucket + + async def test_combined_aggregations(self, agg_index): + index = agg_index["index"] + result = await index.aggregate( + aggregations={ + "avg_price": {"$avg": {"field": "price"}}, + "max_price": {"$max": {"field": "price"}}, + "doc_count": {"$count": {"field": "price"}}, + } + ) + + assert "avg_price" in result + assert "max_price" in result + assert "doc_count" in result + + async def test_aggregate_with_filter(self, agg_index): + index = agg_index["index"] + result = await index.aggregate( + filter={"category": {"$eq": "electronics"}}, + aggregations={"avg_price": {"$avg": {"field": "price"}}}, + ) + + assert "avg_price" in result + assert "value" in result["avg_price"] diff --git a/tests/commands/asyncio/search/test_search_alias_async.py b/tests/commands/asyncio/search/test_search_alias_async.py new file mode 100644 index 0000000..32c58ec --- /dev/null +++ b/tests/commands/asyncio/search/test_search_alias_async.py @@ -0,0 +1,136 @@ +"""Async tests for search alias functionality.""" + +import random +import string + +import pytest +import pytest_asyncio + +from upstash_redis.asyncio import Redis + + +def random_id() -> str: + return "".join(random.choices(string.ascii_lowercase + string.digits, k=8)) + + +@pytest_asyncio.fixture +async def async_redis_client(): + return Redis.from_env() + + +@pytest_asyncio.fixture +async def cleanup_indexes(): + indexes = [] + yield indexes + redis = Redis.from_env() + for index_name in indexes: + try: + index = redis.search.index(index_name) + await index.drop() + except Exception: + pass + + +@pytest.mark.asyncio +class TestAsyncAlias: + """Async tests for alias commands.""" + + async def test_add_alias_via_index( + self, async_redis_client: Redis, cleanup_indexes + ): + name = f"test-async-alias-idx-{random_id()}" + alias_name = f"alias-{random_id()}" + cleanup_indexes.append(name) + + index = await async_redis_client.search.create_index( + name=name, schema={"name": "TEXT"}, data_type="string", prefixes=f"{name}:" + ) + + result = await index.add_alias(alias=alias_name) + assert result == 1 + + await async_redis_client.search.alias.delete(alias=alias_name) + + async def test_alias_add(self, async_redis_client: Redis, cleanup_indexes): + name = f"test-async-alias-add-{random_id()}" + alias_name = f"alias-{random_id()}" + cleanup_indexes.append(name) + + await async_redis_client.search.create_index( + name=name, schema={"name": "TEXT"}, data_type="string", prefixes=f"{name}:" + ) + + result = await async_redis_client.search.alias.add( + index_name=name, alias=alias_name + ) + assert result == 1 + + await async_redis_client.search.alias.delete(alias=alias_name) + + async def test_alias_delete(self, async_redis_client: Redis, cleanup_indexes): + name = f"test-async-alias-del-{random_id()}" + alias_name = f"alias-{random_id()}" + cleanup_indexes.append(name) + + await async_redis_client.search.create_index( + name=name, schema={"name": "TEXT"}, data_type="string", prefixes=f"{name}:" + ) + + await async_redis_client.search.alias.add(index_name=name, alias=alias_name) + + result = await async_redis_client.search.alias.delete(alias=alias_name) + assert result == 1 + + async def test_alias_list(self, async_redis_client: Redis, cleanup_indexes): + name = f"test-async-alias-list-{random_id()}" + alias_name = f"alias-{random_id()}" + cleanup_indexes.append(name) + + await async_redis_client.search.create_index( + name=name, schema={"name": "TEXT"}, data_type="string", prefixes=f"{name}:" + ) + + await async_redis_client.search.alias.add(index_name=name, alias=alias_name) + + aliases = await async_redis_client.search.alias.list() + assert isinstance(aliases, dict) + assert alias_name in aliases + assert aliases[alias_name] == name + + await async_redis_client.search.alias.delete(alias=alias_name) + + async def test_update_alias_to_new_index( + self, async_redis_client: Redis, cleanup_indexes + ): + name1 = f"test-async-alias-upd1-{random_id()}" + name2 = f"test-async-alias-upd2-{random_id()}" + alias_name = f"alias-{random_id()}" + cleanup_indexes.extend([name1, name2]) + + await async_redis_client.search.create_index( + name=name1, + schema={"name": "TEXT"}, + data_type="string", + prefixes=f"{name1}:", + ) + await async_redis_client.search.create_index( + name=name2, + schema={"name": "TEXT"}, + data_type="string", + prefixes=f"{name2}:", + ) + + result1 = await async_redis_client.search.alias.add( + index_name=name1, alias=alias_name + ) + assert result1 == 1 + + result2 = await async_redis_client.search.alias.add( + index_name=name2, alias=alias_name + ) + assert result2 == 2 + + aliases = await async_redis_client.search.alias.list() + assert aliases[alias_name] == name2 + + await async_redis_client.search.alias.delete(alias=alias_name) diff --git a/tests/commands/asyncio/search/test_search_field_types_async.py b/tests/commands/asyncio/search/test_search_field_types_async.py new file mode 100644 index 0000000..16a4e85 --- /dev/null +++ b/tests/commands/asyncio/search/test_search_field_types_async.py @@ -0,0 +1,161 @@ +"""Async tests for KEYWORD and FACET field types.""" + +import json +import random +import string + +import pytest +import pytest_asyncio + +from upstash_redis.asyncio import Redis + + +def random_id() -> str: + return "".join(random.choices(string.ascii_lowercase + string.digits, k=8)) + + +@pytest_asyncio.fixture +async def async_redis_client(): + return Redis.from_env() + + +@pytest.mark.asyncio +class TestAsyncKeywordField: + """Async tests for KEYWORD field type.""" + + @pytest_asyncio.fixture + async def keyword_index(self, async_redis_client): + name = f"test-async-keyword-{random_id()}" + prefix = f"{name}:" + keys = [] + + redis = async_redis_client + + schema = { + "name": "TEXT", + "status": "KEYWORD", + "priority": {"type": "U64", "fast": True}, + } + + index = await redis.search.create_index( + name=name, schema=schema, data_type="string", prefixes=prefix + ) + + test_data = [ + {"name": "Task A", "status": "open", "priority": 1}, + {"name": "Task B", "status": "closed", "priority": 2}, + {"name": "Task C", "status": "open", "priority": 3}, + {"name": "Task D", "status": "in_progress", "priority": 4}, + {"name": "Task E", "status": "closed", "priority": 5}, + ] + + for i, datum in enumerate(test_data): + key = f"{prefix}{i}" + keys.append(key) + await redis.set(key, json.dumps(datum)) + + await index.wait_indexing() + + yield {"index": index, "redis": redis, "keys": keys, "name": name} + + try: + await index.drop() + except Exception: + pass + if keys: + await redis.delete(*keys) + + async def test_keyword_eq(self, keyword_index): + index = keyword_index["index"] + result = await index.query(filter={"status": {"$eq": "open"}}) + assert len(result) == 2 + + async def test_keyword_in(self, keyword_index): + index = keyword_index["index"] + result = await index.query(filter={"status": {"$in": ["open", "closed"]}}) + assert len(result) == 4 + + async def test_keyword_gt(self, keyword_index): + index = keyword_index["index"] + result = await index.query(filter={"status": {"$gt": "in_progress"}}) + assert len(result) >= 1 + + async def test_keyword_gte(self, keyword_index): + index = keyword_index["index"] + result = await index.query(filter={"status": {"$gte": "open"}}) + assert len(result) >= 1 + + async def test_keyword_lt(self, keyword_index): + index = keyword_index["index"] + result = await index.query(filter={"status": {"$lt": "open"}}) + assert len(result) >= 1 + + async def test_keyword_lte(self, keyword_index): + index = keyword_index["index"] + result = await index.query(filter={"status": {"$lte": "open"}}) + assert len(result) >= 1 + + +@pytest.mark.asyncio +class TestAsyncFacetField: + """Async tests for FACET field type.""" + + @pytest_asyncio.fixture + async def facet_index(self, async_redis_client): + name = f"test-async-facet-{random_id()}" + prefix = f"{name}:" + keys = [] + + redis = async_redis_client + + schema = { + "name": "TEXT", + "category": "FACET", + "price": {"type": "F64", "fast": True}, + } + + index = await redis.search.create_index( + name=name, schema=schema, data_type="string", prefixes=prefix + ) + + test_data = [ + {"name": "Laptop", "category": "/electronics/computers", "price": 999.99}, + {"name": "Mouse", "category": "/electronics/peripherals", "price": 29.99}, + {"name": "Book", "category": "/books/fiction", "price": 14.99}, + {"name": "Shirt", "category": "/clothing/tops", "price": 39.99}, + { + "name": "Keyboard", + "category": "/electronics/peripherals", + "price": 79.99, + }, + ] + + for i, datum in enumerate(test_data): + key = f"{prefix}{i}" + keys.append(key) + await redis.set(key, json.dumps(datum)) + + await index.wait_indexing() + + yield {"index": index, "redis": redis, "keys": keys, "name": name} + + try: + await index.drop() + except Exception: + pass + if keys: + await redis.delete(*keys) + + async def test_facet_eq(self, facet_index): + index = facet_index["index"] + result = await index.query( + filter={"category": {"$eq": "/electronics/peripherals"}} + ) + assert len(result) == 2 + + async def test_facet_in(self, facet_index): + index = facet_index["index"] + result = await index.query( + filter={"category": {"$in": ["/electronics/peripherals", "/books/fiction"]}} + ) + assert len(result) == 3 diff --git a/tests/commands/asyncio/search/test_search_scorefunc_async.py b/tests/commands/asyncio/search/test_search_scorefunc_async.py index 56e995c..d6f5323 100644 --- a/tests/commands/asyncio/search/test_search_scorefunc_async.py +++ b/tests/commands/asyncio/search/test_search_scorefunc_async.py @@ -1,9 +1,12 @@ """Tests for async score function queries.""" +from typing import AsyncGenerator + import pytest import pytest_asyncio from upstash_redis.asyncio import Redis as AsyncRedis +from upstash_redis.search import CombineMode, ScoreMode, ScoreModifier def random_id() -> str: @@ -47,7 +50,7 @@ class TestAsyncScoreFunctionQueries: """Tests for async querying with score functions.""" @pytest_asyncio.fixture(scope="class") - async def scorefunc_index(self) -> dict: + async def scorefunc_index(self) -> AsyncGenerator[dict, None]: """Create a test index with data for score function querying.""" redis = AsyncRedis.from_env() name = f"test-scorefunc-async-{random_id()}" @@ -115,7 +118,7 @@ async def test_async_query_with_scorefunc_using_modifier( filter={"name": {"$eq": "Laptop"}}, score_func={ "field": "popularity", - "modifier": "LOG1P", + "modifier": ScoreModifier.LOG1P, "factor": 2, }, ) @@ -134,7 +137,7 @@ async def test_async_query_with_scorefunc_using_scoremode_replace( filter={"name": {"$eq": "Laptop"}}, score_func={ "field": "popularity", - "scoreMode": "REPLACE", + "scoreMode": ScoreMode.REPLACE, }, ) @@ -153,10 +156,10 @@ async def test_async_query_with_multiple_field_values_combined( filter={"name": {"$eq": "Laptop"}}, score_func={ "fields": [ - {"field": "popularity", "modifier": "LOG1P"}, - {"field": "recency", "modifier": "LOG1P"}, + {"field": "popularity", "modifier": ScoreModifier.LOG1P}, + {"field": "recency", "modifier": ScoreModifier.LOG1P}, ], - "combineMode": "SUM", + "combineMode": CombineMode.SUM, }, ) diff --git a/tests/commands/search/test_search_aggregate.py b/tests/commands/search/test_search_aggregate.py new file mode 100644 index 0000000..8ef3f0a --- /dev/null +++ b/tests/commands/search/test_search_aggregate.py @@ -0,0 +1,317 @@ +"""Tests for search aggregation functionality.""" + +import json +from typing import Generator + +import pytest + +from upstash_redis import Redis + + +def random_id() -> str: + import random + import string + + return "".join(random.choices(string.ascii_lowercase + string.digits, k=8)) + + +@pytest.fixture +def redis_client() -> Redis: + return Redis.from_env() + + +class TestAggregation: + """Tests for aggregation queries.""" + + @pytest.fixture(scope="class") + def agg_index(self) -> Generator[dict, None, None]: + name = f"test-agg-{random_id()}" + prefix = f"{name}:" + keys = [] + + redis = Redis.from_env() + + schema = { + "name": "TEXT", + "category": "KEYWORD", + "price": {"type": "F64", "fast": True}, + "stock": {"type": "U64", "fast": True}, + "active": "BOOL", + } + + index = redis.search.create_index( + name=name, schema=schema, data_type="string", prefixes=prefix + ) + + test_data = [ + { + "name": "Laptop Pro", + "category": "electronics", + "price": 1299.99, + "stock": 50, + "active": True, + }, + { + "name": "Laptop Basic", + "category": "electronics", + "price": 599.99, + "stock": 100, + "active": True, + }, + { + "name": "Wireless Mouse", + "category": "electronics", + "price": 29.99, + "stock": 200, + "active": True, + }, + { + "name": "USB Cable", + "category": "accessories", + "price": 9.99, + "stock": 500, + "active": True, + }, + { + "name": "Phone Case", + "category": "accessories", + "price": 19.99, + "stock": 300, + "active": False, + }, + { + "name": "Keyboard", + "category": "electronics", + "price": 79.99, + "stock": 150, + "active": True, + }, + { + "name": "Monitor", + "category": "electronics", + "price": 399.99, + "stock": 75, + "active": True, + }, + { + "name": "Headphones", + "category": "audio", + "price": 149.99, + "stock": 120, + "active": True, + }, + { + "name": "Speaker", + "category": "audio", + "price": 59.99, + "stock": 80, + "active": True, + }, + ] + + for i, datum in enumerate(test_data): + key = f"{prefix}{i}" + keys.append(key) + redis.set(key, json.dumps(datum)) + + index.wait_indexing() + + yield {"index": index, "redis": redis, "keys": keys, "name": name} + + try: + index.drop() + except Exception: + pass + if keys: + redis.delete(*keys) + + def test_avg_aggregation(self, agg_index): + index = agg_index["index"] + result = index.aggregate( + aggregations={"avg_price": {"$avg": {"field": "price"}}} + ) + + assert "avg_price" in result + assert "value" in result["avg_price"] + assert isinstance(result["avg_price"]["value"], (int, float)) + + def test_sum_aggregation(self, agg_index): + index = agg_index["index"] + result = index.aggregate( + aggregations={"total_price": {"$sum": {"field": "price"}}} + ) + + assert "total_price" in result + assert "value" in result["total_price"] + + def test_min_aggregation(self, agg_index): + index = agg_index["index"] + result = index.aggregate( + aggregations={"min_price": {"$min": {"field": "price"}}} + ) + + assert "min_price" in result + assert "value" in result["min_price"] + + def test_max_aggregation(self, agg_index): + index = agg_index["index"] + result = index.aggregate( + aggregations={"max_price": {"$max": {"field": "price"}}} + ) + + assert "max_price" in result + assert "value" in result["max_price"] + + def test_count_aggregation(self, agg_index): + index = agg_index["index"] + result = index.aggregate( + aggregations={"doc_count": {"$count": {"field": "price"}}} + ) + + assert "doc_count" in result + assert "value" in result["doc_count"] + + def test_cardinality_aggregation(self, agg_index): + index = agg_index["index"] + result = index.aggregate( + aggregations={"unique_stock": {"$cardinality": {"field": "stock"}}} + ) + + assert "unique_stock" in result + assert "value" in result["unique_stock"] + + def test_stats_aggregation(self, agg_index): + index = agg_index["index"] + result = index.aggregate( + aggregations={"price_stats": {"$stats": {"field": "price"}}} + ) + + assert "price_stats" in result + stats = result["price_stats"] + assert "count" in stats + assert "min" in stats + assert "max" in stats + assert "sum" in stats + assert "avg" in stats + + def test_extended_stats_aggregation(self, agg_index): + index = agg_index["index"] + result = index.aggregate( + aggregations={"price_estats": {"$extendedStats": {"field": "price"}}} + ) + + assert "price_estats" in result + estats = result["price_estats"] + assert "count" in estats + assert "min" in estats + assert "max" in estats + + def test_percentiles_aggregation(self, agg_index): + index = agg_index["index"] + result = index.aggregate( + aggregations={ + "price_pct": { + "$percentiles": { + "field": "price", + "percents": [50, 95], + } + } + } + ) + + assert "price_pct" in result + + def test_terms_aggregation(self, agg_index): + index = agg_index["index"] + result = index.aggregate( + aggregations={"by_category": {"$terms": {"field": "category"}}} + ) + + assert "by_category" in result + assert "buckets" in result["by_category"] + buckets = result["by_category"]["buckets"] + assert len(buckets) > 0 + for bucket in buckets: + assert "key" in bucket + assert "docCount" in bucket + + def test_range_aggregation(self, agg_index): + index = agg_index["index"] + result = index.aggregate( + aggregations={ + "price_ranges": { + "$range": { + "field": "price", + "ranges": [ + {"to": 50}, + {"from": 50, "to": 500}, + {"from": 500}, + ], + } + } + } + ) + + assert "price_ranges" in result + assert "buckets" in result["price_ranges"] + + def test_histogram_aggregation(self, agg_index): + index = agg_index["index"] + result = index.aggregate( + aggregations={ + "price_hist": { + "$histogram": { + "field": "price", + "interval": 200, + } + } + } + ) + + assert "price_hist" in result + assert "buckets" in result["price_hist"] + + def test_nested_sub_aggregation(self, agg_index): + index = agg_index["index"] + result = index.aggregate( + aggregations={ + "by_category": { + "$terms": { + "field": "category", + }, + "$aggs": { + "avg_price": {"$avg": {"field": "price"}}, + }, + } + } + ) + + assert "by_category" in result + assert "buckets" in result["by_category"] + for bucket in result["by_category"]["buckets"]: + assert "key" in bucket + assert "avg_price" in bucket + + def test_combined_aggregations(self, agg_index): + index = agg_index["index"] + result = index.aggregate( + aggregations={ + "avg_price": {"$avg": {"field": "price"}}, + "max_price": {"$max": {"field": "price"}}, + "doc_count": {"$count": {"field": "price"}}, + } + ) + + assert "avg_price" in result + assert "max_price" in result + assert "doc_count" in result + + def test_aggregate_with_filter(self, agg_index): + index = agg_index["index"] + result = index.aggregate( + filter={"category": {"$eq": "electronics"}}, + aggregations={"avg_price": {"$avg": {"field": "price"}}}, + ) + + assert "avg_price" in result + assert "value" in result["avg_price"] diff --git a/tests/commands/search/test_search_alias.py b/tests/commands/search/test_search_alias.py new file mode 100644 index 0000000..ce60814 --- /dev/null +++ b/tests/commands/search/test_search_alias.py @@ -0,0 +1,134 @@ +"""Tests for search alias functionality.""" + +import pytest + +from upstash_redis import Redis + + +def random_id() -> str: + import random + import string + + return "".join(random.choices(string.ascii_lowercase + string.digits, k=8)) + + +@pytest.fixture +def redis_client() -> Redis: + return Redis.from_env() + + +@pytest.fixture +def cleanup_indexes(): + indexes = [] + yield indexes + redis = Redis.from_env() + for index_name in indexes: + try: + redis.search.index(index_name).drop() + except Exception: + pass + + +class TestAlias: + """Tests for alias commands.""" + + def test_add_alias_via_index(self, redis_client: Redis, cleanup_indexes): + """Test adding an alias via the index method.""" + name = f"test-alias-idx-{random_id()}" + alias_name = f"alias-{random_id()}" + cleanup_indexes.append(name) + + index = redis_client.search.create_index( + name=name, schema={"name": "TEXT"}, data_type="string", prefixes=f"{name}:" + ) + + result = index.add_alias(alias=alias_name) + assert result == 1 + + # Clean up alias + redis_client.search.alias.delete(alias=alias_name) + + def test_alias_add(self, redis_client: Redis, cleanup_indexes): + """Test adding an alias via search.alias.add.""" + name = f"test-alias-add-{random_id()}" + alias_name = f"alias-{random_id()}" + cleanup_indexes.append(name) + + redis_client.search.create_index( + name=name, schema={"name": "TEXT"}, data_type="string", prefixes=f"{name}:" + ) + + result = redis_client.search.alias.add(index_name=name, alias=alias_name) + assert result == 1 + + # Clean up alias + redis_client.search.alias.delete(alias=alias_name) + + def test_alias_delete(self, redis_client: Redis, cleanup_indexes): + """Test deleting an alias.""" + name = f"test-alias-del-{random_id()}" + alias_name = f"alias-{random_id()}" + cleanup_indexes.append(name) + + redis_client.search.create_index( + name=name, schema={"name": "TEXT"}, data_type="string", prefixes=f"{name}:" + ) + + redis_client.search.alias.add(index_name=name, alias=alias_name) + + result = redis_client.search.alias.delete(alias=alias_name) + assert result == 1 + + def test_alias_list(self, redis_client: Redis, cleanup_indexes): + """Test listing aliases.""" + name = f"test-alias-list-{random_id()}" + alias_name = f"alias-{random_id()}" + cleanup_indexes.append(name) + + redis_client.search.create_index( + name=name, schema={"name": "TEXT"}, data_type="string", prefixes=f"{name}:" + ) + + redis_client.search.alias.add(index_name=name, alias=alias_name) + + aliases = redis_client.search.alias.list() + assert isinstance(aliases, dict) + assert alias_name in aliases + assert aliases[alias_name] == name + + # Clean up alias + redis_client.search.alias.delete(alias=alias_name) + + def test_update_alias_to_new_index(self, redis_client: Redis, cleanup_indexes): + """Test updating an existing alias to point to a new index.""" + name1 = f"test-alias-upd1-{random_id()}" + name2 = f"test-alias-upd2-{random_id()}" + alias_name = f"alias-{random_id()}" + cleanup_indexes.extend([name1, name2]) + + redis_client.search.create_index( + name=name1, + schema={"name": "TEXT"}, + data_type="string", + prefixes=f"{name1}:", + ) + redis_client.search.create_index( + name=name2, + schema={"name": "TEXT"}, + data_type="string", + prefixes=f"{name2}:", + ) + + # Add alias pointing to first index + result1 = redis_client.search.alias.add(index_name=name1, alias=alias_name) + assert result1 == 1 + + # Update alias to point to second index + result2 = redis_client.search.alias.add(index_name=name2, alias=alias_name) + assert result2 == 2 + + aliases = redis_client.search.alias.list() + assert aliases[alias_name] == name2 + + # Clean up alias + redis_client.search.alias.delete(alias=alias_name) diff --git a/tests/commands/search/test_search_field_types.py b/tests/commands/search/test_search_field_types.py new file mode 100644 index 0000000..430fe30 --- /dev/null +++ b/tests/commands/search/test_search_field_types.py @@ -0,0 +1,158 @@ +"""Tests for KEYWORD and FACET field types.""" + +import json +from typing import Generator + +import pytest + +from upstash_redis import Redis + + +def random_id() -> str: + import random + import string + + return "".join(random.choices(string.ascii_lowercase + string.digits, k=8)) + + +@pytest.fixture +def redis_client() -> Redis: + return Redis.from_env() + + +class TestKeywordField: + """Tests for KEYWORD field type.""" + + @pytest.fixture(scope="class") + def keyword_index(self) -> Generator[dict, None, None]: + name = f"test-keyword-{random_id()}" + prefix = f"{name}:" + keys = [] + + redis = Redis.from_env() + + schema = { + "name": "TEXT", + "status": "KEYWORD", + "priority": {"type": "U64", "fast": True}, + } + + index = redis.search.create_index( + name=name, schema=schema, data_type="string", prefixes=prefix + ) + + test_data = [ + {"name": "Task A", "status": "open", "priority": 1}, + {"name": "Task B", "status": "closed", "priority": 2}, + {"name": "Task C", "status": "open", "priority": 3}, + {"name": "Task D", "status": "in_progress", "priority": 4}, + {"name": "Task E", "status": "closed", "priority": 5}, + ] + + for i, datum in enumerate(test_data): + key = f"{prefix}{i}" + keys.append(key) + redis.set(key, json.dumps(datum)) + + index.wait_indexing() + + yield {"index": index, "redis": redis, "keys": keys, "name": name} + + try: + index.drop() + except Exception: + pass + if keys: + redis.delete(*keys) + + def test_keyword_eq(self, keyword_index): + index = keyword_index["index"] + result = index.query(filter={"status": {"$eq": "open"}}) + assert len(result) == 2 + + def test_keyword_in(self, keyword_index): + index = keyword_index["index"] + result = index.query(filter={"status": {"$in": ["open", "closed"]}}) + assert len(result) == 4 + + def test_keyword_gt(self, keyword_index): + index = keyword_index["index"] + result = index.query(filter={"status": {"$gt": "in_progress"}}) + assert len(result) >= 1 + + def test_keyword_gte(self, keyword_index): + index = keyword_index["index"] + result = index.query(filter={"status": {"$gte": "open"}}) + assert len(result) >= 1 + + def test_keyword_lt(self, keyword_index): + index = keyword_index["index"] + result = index.query(filter={"status": {"$lt": "open"}}) + assert len(result) >= 1 + + def test_keyword_lte(self, keyword_index): + index = keyword_index["index"] + result = index.query(filter={"status": {"$lte": "open"}}) + assert len(result) >= 1 + + +class TestFacetField: + """Tests for FACET field type.""" + + @pytest.fixture(scope="class") + def facet_index(self) -> Generator[dict, None, None]: + name = f"test-facet-{random_id()}" + prefix = f"{name}:" + keys = [] + + redis = Redis.from_env() + + schema = { + "name": "TEXT", + "category": "FACET", + "price": {"type": "F64", "fast": True}, + } + + index = redis.search.create_index( + name=name, schema=schema, data_type="string", prefixes=prefix + ) + + test_data = [ + {"name": "Laptop", "category": "/electronics/computers", "price": 999.99}, + {"name": "Mouse", "category": "/electronics/peripherals", "price": 29.99}, + {"name": "Book", "category": "/books/fiction", "price": 14.99}, + {"name": "Shirt", "category": "/clothing/tops", "price": 39.99}, + { + "name": "Keyboard", + "category": "/electronics/peripherals", + "price": 79.99, + }, + ] + + for i, datum in enumerate(test_data): + key = f"{prefix}{i}" + keys.append(key) + redis.set(key, json.dumps(datum)) + + index.wait_indexing() + + yield {"index": index, "redis": redis, "keys": keys, "name": name} + + try: + index.drop() + except Exception: + pass + if keys: + redis.delete(*keys) + + def test_facet_eq(self, facet_index): + index = facet_index["index"] + result = index.query(filter={"category": {"$eq": "/electronics/peripherals"}}) + assert len(result) == 2 + + def test_facet_in(self, facet_index): + index = facet_index["index"] + result = index.query( + filter={"category": {"$in": ["/electronics/peripherals", "/books/fiction"]}} + ) + assert len(result) == 3 diff --git a/tests/commands/search/test_search_scorefunc.py b/tests/commands/search/test_search_scorefunc.py index 16eaf75..2f2fb00 100644 --- a/tests/commands/search/test_search_scorefunc.py +++ b/tests/commands/search/test_search_scorefunc.py @@ -7,6 +7,7 @@ from upstash_redis import Redis from upstash_redis.commands import SearchIndexCommands +from upstash_redis.search import CombineMode, ScoreMode, ScoreModifier class ScoreFuncFixture(TypedDict): @@ -103,7 +104,7 @@ def test_query_with_simple_scorefunc(self, scorefunc_index: ScoreFuncFixture): index = scorefunc_index["index"] result = index.query( filter={"name": {"$eq": "Laptop"}}, - score_func={"field": "popularity", "scoreMode": "REPLACE"}, + score_func={"field": "popularity", "scoreMode": ScoreMode.REPLACE}, ) assert len(result) == 4 @@ -113,16 +114,18 @@ def test_query_with_simple_scorefunc(self, scorefunc_index: ScoreFuncFixture): assert scores_by_name["Laptop Pro"] == pytest.approx(1000.0) assert scores_by_name["Laptop Basic"] == pytest.approx(500.0) - def test_query_with_scorefunc_using_modifier(self, scorefunc_index: ScoreFuncFixture): + def test_query_with_scorefunc_using_modifier( + self, scorefunc_index: ScoreFuncFixture + ): """Test querying with scoreFunc using modifier.""" index = scorefunc_index["index"] result = index.query( filter={"name": {"$eq": "Laptop"}}, score_func={ "field": "popularity", - "modifier": "LOG1P", + "modifier": ScoreModifier.LOG1P, "factor": 2, - "scoreMode": "REPLACE", + "scoreMode": ScoreMode.REPLACE, }, ) @@ -141,7 +144,7 @@ def test_query_with_scorefunc_using_scoremode_replace( filter={"name": {"$eq": "Laptop"}}, score_func={ "field": "popularity", - "scoreMode": "REPLACE", + "scoreMode": ScoreMode.REPLACE, }, ) @@ -151,18 +154,20 @@ def test_query_with_scorefunc_using_scoremode_replace( assert scores_by_name["Laptop Pro"] == pytest.approx(1000.0) assert scores_by_name["Laptop Basic"] == pytest.approx(500.0) - def test_query_with_multiple_field_values_combined(self, scorefunc_index: ScoreFuncFixture): + def test_query_with_multiple_field_values_combined( + self, scorefunc_index: ScoreFuncFixture + ): """Test querying with multiple field values combined.""" index = scorefunc_index["index"] result = index.query( filter={"name": {"$eq": "Laptop"}}, score_func={ "fields": [ - {"field": "popularity", "modifier": "LOG1P"}, - {"field": "recency", "modifier": "LOG1P"}, + {"field": "popularity", "modifier": ScoreModifier.LOG1P}, + {"field": "recency", "modifier": ScoreModifier.LOG1P}, ], - "combineMode": "SUM", - "scoreMode": "REPLACE", + "combineMode": CombineMode.SUM, + "scoreMode": ScoreMode.REPLACE, }, ) @@ -178,16 +183,18 @@ def test_query_with_multiple_field_values_combined(self, scorefunc_index: ScoreF math.log10(1 + 500) + math.log10(1 + 200) ) - def test_query_with_scorefunc_modifier_and_missing(self, scorefunc_index: ScoreFuncFixture): + def test_query_with_scorefunc_modifier_and_missing( + self, scorefunc_index: ScoreFuncFixture + ): """Test querying with scoreFunc using modifier and missing value.""" index = scorefunc_index["index"] result = index.query( filter={"name": {"$eq": "Laptop"}}, score_func={ "field": "popularity", - "modifier": "LOG1P", + "modifier": ScoreModifier.LOG1P, "missing": 1, - "scoreMode": "REPLACE", + "scoreMode": ScoreMode.REPLACE, }, ) @@ -195,15 +202,17 @@ def test_query_with_scorefunc_modifier_and_missing(self, scorefunc_index: ScoreF scores_by_name = {r.data["name"]: r.score for r in result if r.data} assert scores_by_name["Laptop Missing"] == pytest.approx(math.log10(1 + 1)) - def test_query_with_multiple_fields_and_scoremode(self, scorefunc_index: ScoreFuncFixture): + def test_query_with_multiple_fields_and_scoremode( + self, scorefunc_index: ScoreFuncFixture + ): """Test querying with multiple field values and scoreMode.""" index = scorefunc_index["index"] result = index.query( filter={"name": {"$eq": "Laptop"}}, score_func={ "fields": ["popularity", "recency"], - "combineMode": "MULTIPLY", - "scoreMode": "REPLACE", + "combineMode": CombineMode.MULTIPLY, + "scoreMode": ScoreMode.REPLACE, }, ) diff --git a/upstash_redis/commands.py b/upstash_redis/commands.py index 2313b2a..fbe8743 100644 --- a/upstash_redis/commands.py +++ b/upstash_redis/commands.py @@ -5577,6 +5577,45 @@ def index(self, name: str) -> "SearchIndexCommands": """Initialize a SearchIndexCommands instance for an existing index.""" return SearchIndexCommands(self.client, name) + @property + def alias(self) -> "SearchAliasCommands": + """Access alias commands.""" + return SearchAliasCommands(self.client) + + +class SearchAliasCommands: + """Commands for managing search index aliases.""" + + def __init__(self, client: Commands): + self.client = client + + def add(self, *, index_name: str, alias: str) -> ResponseT: + """ + Add or update an alias for an index. + + Returns 1 if alias was created, 2 if updated. + """ + command: List = ["SEARCH.ALIASADD", alias, index_name] + return self.client.execute(command) + + def delete(self, *, alias: str) -> ResponseT: + """ + Delete an alias. + + Returns 1 if alias was deleted, 0 if not found. + """ + command: List = ["SEARCH.ALIASDEL", alias] + return self.client.execute(command) + + def list(self) -> ResponseT: + """ + List all aliases. + + Returns a dict mapping alias names to index names. + """ + command: List = ["SEARCH.LISTALIASES"] + return self.client.execute(command) + class SearchIndexCommands: """ @@ -5681,6 +5720,40 @@ def describe(self) -> ResponseT: command = ["SEARCH.DESCRIBE", self.name] return self.client.execute(command) + def aggregate( + self, + *, + filter: Optional[Dict[str, Any]] = None, + aggregations: Dict[str, Any], + ) -> ResponseT: + """ + Run aggregation queries on the index. + + Example: + ```python + result = index.aggregate( + filter={"category": {"$eq": "electronics"}}, + aggregations={"avg_price": {"$avg": {"field": "price"}}} + ) + ``` + """ + command: List = [ + "SEARCH.AGGREGATE", + self.name, + json.dumps(filter or {}), + json.dumps(aggregations), + ] + return self.client.execute(command) + + def add_alias(self, *, alias: str) -> ResponseT: + """ + Add or update an alias for this index. + + Returns 1 if alias was created, 2 if updated. + """ + command: List = ["SEARCH.ALIASADD", alias, self.name] + return self.client.execute(command) + def drop(self) -> ResponseT: """ Drop (delete) the index. @@ -5786,6 +5859,7 @@ def execute(self) -> ResponseT: AsyncJsonCommands = JsonCommands AsyncSearchCommands = SearchCommands AsyncSearchIndexCommands = SearchIndexCommands +AsyncSearchAliasCommands = SearchAliasCommands AsyncBitFieldCommands = BitFieldCommands AsyncBitFieldROCommands = BitFieldROCommands PipelineCommands = Commands diff --git a/upstash_redis/commands.pyi b/upstash_redis/commands.pyi index 335606f..0b8a85c 100644 --- a/upstash_redis/commands.pyi +++ b/upstash_redis/commands.pyi @@ -2040,6 +2040,12 @@ class AsyncJsonCommands: async def toggle(self, key: str, path: str = "$") -> List[Union[bool, None]]: ... async def type(self, key: str, path: str = "$") -> List[str]: ... +class SearchAliasCommands: + def __init__(self, client: Commands): ... + def add(self, *, index_name: str, alias: str) -> int: ... + def delete(self, *, alias: str) -> int: ... + def list(self) -> Dict[str, str]: ... + class SearchCommands: def __init__(self, client: Commands): ... def create_index( @@ -2057,11 +2063,13 @@ class SearchCommands: self, name: str, ) -> SearchIndexCommands: ... + @property + def alias(self) -> SearchAliasCommands: ... class SearchIndexCommands: def __init__(self, client: Commands, name: str): ... - def wait_indexing(self) -> None: ... - def describe(self) -> IndexDescription: ... + def wait_indexing(self) -> int: ... + def describe(self) -> Optional[IndexDescription]: ... def query( self, *, @@ -2074,7 +2082,20 @@ class SearchIndexCommands: score_func: Optional[ScoreFunc] = None, ) -> List[QueryResult]: ... def count(self, *, filter: Dict[str, Any]) -> CountResult: ... - def drop(self) -> None: ... + def aggregate( + self, + *, + filter: Optional[Dict[str, Any]] = None, + aggregations: Dict[str, Any], + ) -> Dict[str, Any]: ... + def add_alias(self, *, alias: str) -> int: ... + def drop(self) -> int: ... + +class AsyncSearchAliasCommands: + def __init__(self, client: AsyncCommands): ... + async def add(self, *, index_name: str, alias: str) -> int: ... + async def delete(self, *, alias: str) -> int: ... + async def list(self) -> Dict[str, str]: ... class AsyncSearchCommands: def __init__(self, client: AsyncCommands): ... @@ -2090,11 +2111,13 @@ class AsyncSearchCommands: exists_ok: bool = False, ) -> AsyncSearchIndexCommands: ... def index(self, name: str) -> AsyncSearchIndexCommands: ... + @property + def alias(self) -> AsyncSearchAliasCommands: ... class AsyncSearchIndexCommands: def __init__(self, client: AsyncCommands, name: str): ... - async def wait_indexing(self) -> None: ... - async def describe(self) -> IndexDescription: ... + async def wait_indexing(self) -> int: ... + async def describe(self) -> Optional[IndexDescription]: ... async def query( self, *, @@ -2107,4 +2130,11 @@ class AsyncSearchIndexCommands: score_func: Optional[ScoreFunc] = None, ) -> List[QueryResult]: ... async def count(self, *, filter: Dict[str, Any]) -> CountResult: ... - async def drop(self) -> None: ... + async def aggregate( + self, + *, + filter: Optional[Dict[str, Any]] = None, + aggregations: Dict[str, Any], + ) -> Dict[str, Any]: ... + async def add_alias(self, *, alias: str) -> int: ... + async def drop(self) -> int: ... diff --git a/upstash_redis/format.py b/upstash_redis/format.py index ad1e04c..c3d1489 100644 --- a/upstash_redis/format.py +++ b/upstash_redis/format.py @@ -3,6 +3,7 @@ from upstash_redis.commands import SearchIndexCommands from upstash_redis.search import ( + deserialize_aggregate_response, deserialize_describe_response, deserialize_query_response, CountResult, @@ -232,13 +233,17 @@ def format_xpending_response(res, command, _): return res -def format_search_query_response(res: List[Any], _, __): +def format_search_query_response(res: Any, _, __): """Format SEARCH.QUERY response into structured results.""" + if not isinstance(res, list): + return [] return deserialize_query_response(res) -def format_search_describe_response(res: List[Any], _, __): +def format_search_describe_response(res: Any, _, __): """Format SEARCH.DESCRIBE response into index description.""" + if not isinstance(res, list) or len(res) == 0: + return None return deserialize_describe_response(res) @@ -258,6 +263,35 @@ def format_search_create_response(res: Any, command: List[str], client: Any): return SearchIndexCommands(client, index_name) +def format_search_aggregate_response(res: Any, _, __): + """Format SEARCH.AGGREGATE response into parsed aggregation results.""" + return deserialize_aggregate_response(res) + + +def format_search_alias_add_response(res: Any, _, __): + """Format SEARCH.ALIASADD response. Returns 1 (created) or 2 (updated).""" + return int(res) + + +def format_search_alias_del_response(res: Any, _, __): + """Format SEARCH.ALIASDEL response. Returns 1 (deleted) or 0 (not found).""" + return int(res) + + +def format_search_list_aliases_response(res: Any, _, __): + """Format SEARCH.LISTALIASES response into {alias: index_name} dict.""" + if res == 0 or (isinstance(res, list) and len(res) == 0): + return {} + if not isinstance(res, list): + return {} + + aliases: Dict[str, str] = {} + for pair in res: + if isinstance(pair, list) and len(pair) == 2: + aliases[pair[0]] = pair[1] + return aliases + + FORMATTERS: Dict[str, Callable] = { "COPY": to_bool, "EXPIRE": to_bool, @@ -337,6 +371,10 @@ def format_search_create_response(res: Any, command: List[str], client: Any): "SEARCH.QUERY": format_search_query_response, "SEARCH.DESCRIBE": format_search_describe_response, "SEARCH.COUNT": format_search_count_response, + "SEARCH.AGGREGATE": format_search_aggregate_response, + "SEARCH.ALIASADD": format_search_alias_add_response, + "SEARCH.ALIASDEL": format_search_alias_del_response, + "SEARCH.LISTALIASES": format_search_list_aliases_response, } diff --git a/upstash_redis/search.py b/upstash_redis/search.py index 27390aa..8d2f733 100644 --- a/upstash_redis/search.py +++ b/upstash_redis/search.py @@ -20,6 +20,8 @@ class FieldType(str, enum.Enum): F64 = "F64" BOOL = "BOOL" DATE = "DATE" + KEYWORD = "KEYWORD" + FACET = "FACET" class FieldOptions(TypedDict, total=False): @@ -75,6 +77,7 @@ class ScoreModifier(str, enum.Enum): SQRT = "SQRT" RECIPROCAL = "RECIPROCAL" + class ScoreMode(str, enum.Enum): """Score combination modes.""" @@ -277,3 +280,106 @@ def deserialize_describe_response(raw_response: List[Any]) -> IndexDescription: language=language, schema=schema, ) + + +def _coerce_numeric_string(value: Any) -> Any: + """Convert numeric strings to numbers.""" + if isinstance(value, str): + try: + if "." in value: + return float(value) + return int(value) + except ValueError: + return value + return value + + +def _parse_stats_value(arr: List[Any]) -> Dict[str, Any]: + """Parse a stats-like flat key-value array into a dict.""" + result: Dict[str, Any] = {} + for i in range(0, len(arr), 2): + key = arr[i] + value = arr[i + 1] + + if isinstance(value, list) and len(value) > 0: + if isinstance(value[0], str): + # Nested stats (e.g. stdDeviationBounds) + result[key] = _parse_stats_value(value) + elif ( + isinstance(value[0], list) + and len(value[0]) > 0 + and isinstance(value[0][0], str) + ): + # Percentiles unkeyed: [[key, val, value, val], ...] + result[key] = [_parse_stats_value(item) for item in value] + else: + result[key] = value + else: + result[key] = _coerce_numeric_string(value) + + return result + + +def _parse_buckets_value(arr: List[Any]) -> Dict[str, Any]: + """Parse a bucket aggregation value starting with 'buckets'.""" + if arr[0] == "buckets" and isinstance(arr[1], list): + buckets = [] + for bucket in arr[1]: + bucket_obj: Dict[str, Any] = {} + for i in range(0, len(bucket), 2): + key = bucket[i] + value = bucket[i + 1] + if ( + isinstance(value, list) + and len(value) > 0 + and isinstance(value[0], str) + ): + bucket_obj[key] = _parse_stats_value(value) + else: + bucket_obj[key] = value + buckets.append(bucket_obj) + + result: Dict[str, Any] = {"buckets": buckets} + # Extra key-value pairs after the buckets array (sumOtherDocCount, etc.) + for i in range(2, len(arr), 2): + result[arr[i]] = _coerce_numeric_string(arr[i + 1]) + return result + + return {"raw": arr} + + +def _parse_aggregation_array(arr: List[Any]) -> Dict[str, Any]: + """Parse a top-level aggregation response array.""" + result: Dict[str, Any] = {} + for i in range(0, len(arr), 2): + key = arr[i] + value = arr[i + 1] + + if isinstance(value, list): + if len(value) > 0 and isinstance(value[0], str): + if value[0] == "buckets": + result[key] = _parse_buckets_value(value) + else: + result[key] = _parse_stats_value(value) + else: + result[key] = _parse_aggregation_array(value) + else: + result[key] = value + + return result + + +def deserialize_aggregate_response(raw_response: Any) -> Dict[str, Any]: + """ + Deserialize raw aggregate response into structured results. + + Args: + raw_response: Raw response from SEARCH.AGGREGATE + + Returns: + Parsed aggregation results + """ + if not isinstance(raw_response, list) or len(raw_response) == 0: + return {} + + return _parse_aggregation_array(raw_response) diff --git a/upstash_redis/utils.py b/upstash_redis/utils.py index aafbbe0..05257ab 100644 --- a/upstash_redis/utils.py +++ b/upstash_redis/utils.py @@ -1,7 +1,7 @@ from dataclasses import dataclass -from typing import Any, List, Literal, Optional, Dict +from typing import Any, Dict, List, Literal, Optional, cast -from upstash_redis.search import ScoreFunc +from upstash_redis.search import ScoreByMultipleFields, ScoreFunc from upstash_redis.typing import FloatMinMaxT @@ -137,25 +137,26 @@ def build_score_func(command: List, score_func: ScoreFunc) -> None: command.extend(("FIELDVALUE", score_func)) elif "fields" in score_func: # Multiple fields - if "combineMode" in score_func: - command.extend(("COMBINEMODE", score_func["combineMode"].upper())) - if "scoreMode" in score_func: - command.extend(("SCOREMODE", score_func["scoreMode"].upper())) + multi = cast(ScoreByMultipleFields, score_func) + if "combineMode" in multi: + command.extend(("COMBINEMODE", multi["combineMode"].upper())) + if "scoreMode" in multi: + command.extend(("SCOREMODE", multi["scoreMode"].upper())) - for field_spec in score_func["fields"]: + for field_spec in multi["fields"]: if isinstance(field_spec, str): command.extend(("FIELDVALUE", field_spec)) else: - build_field_value(command, field_spec) + build_field_value(command, dict(field_spec)) else: # Single field with options - if "scoreMode" in score_func: - command.extend(("SCOREMODE", score_func["scoreMode"].upper())) - build_field_value(command, score_func) + single: Dict[str, Any] = dict(score_func) + if "scoreMode" in single: + command.extend(("SCOREMODE", single["scoreMode"].upper())) + build_field_value(command, single) -def build_field_value( - command: List, field_spec: Dict[str, Any] -) -> None: + +def build_field_value(command: List, field_spec: Dict[str, Any]) -> None: """Build a FIELDVALUE portion with modifiers.""" command.extend(("FIELDVALUE", field_spec["field"])) @@ -164,4 +165,4 @@ def build_field_value( if "factor" in field_spec: command.extend(("FACTOR", field_spec["factor"])) if "missing" in field_spec: - command.extend(("MISSING", field_spec["missing"])) \ No newline at end of file + command.extend(("MISSING", field_spec["missing"])) From dac6c13b99b0c8b97e75ff7e34817bcdbb94b5f5 Mon Sep 17 00:00:00 2001 From: CahidArda Date: Wed, 18 Mar 2026 16:30:30 +0300 Subject: [PATCH 12/12] fix: fix highlight serialization and make query filter optional - HIGHLIGHT FIELDS now sends count prefix + spread fields to match protocol - HIGHLIGHT tags check uses safe key lookup to avoid KeyError - query filter parameter is now optional, defaults to match-all ({}) - index() and alias return types use ResponseT for consistency Co-Authored-By: Claude Opus 4.6 (1M context) --- upstash_redis/commands.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/upstash_redis/commands.py b/upstash_redis/commands.py index fbe8743..3a7fa76 100644 --- a/upstash_redis/commands.py +++ b/upstash_redis/commands.py @@ -5573,12 +5573,12 @@ def create_index( return self.client.execute(command) - def index(self, name: str) -> "SearchIndexCommands": + def index(self, name: str) -> ResponseT: """Initialize a SearchIndexCommands instance for an existing index.""" return SearchIndexCommands(self.client, name) @property - def alias(self) -> "SearchAliasCommands": + def alias(self) -> ResponseT: """Access alias commands.""" return SearchAliasCommands(self.client) @@ -5631,7 +5631,7 @@ def __init__(self, client: Commands, name: str): def query( self, *, - filter: Dict[str, Any], + filter: Optional[Dict[str, Any]] = None, limit: Optional[int] = None, offset: Optional[int] = None, order_by: Optional[Dict[str, Union[Order, str]]] = None, @@ -5647,7 +5647,7 @@ def query( results = index.query(filter={"name": {"$eq": "Laptop"}}) ``` """ - command: List = ["SEARCH.QUERY", self.name, json.dumps(filter)] + command: List = ["SEARCH.QUERY", self.name, json.dumps(filter or {})] if limit: command.extend(("LIMIT", limit)) @@ -5672,8 +5672,9 @@ def query( command.extend(("SELECT", len(selected_fields), *selected_fields)) if highlight: - command.extend(("HIGHLIGHT", "FIELDS", highlight["fields"])) - if highlight["tags"]: + fields = highlight["fields"] + command.extend(("HIGHLIGHT", "FIELDS", len(fields), *fields)) + if "tags" in highlight and highlight["tags"]: open_tag, close_tag = highlight["tags"] command.extend(("TAGS", open_tag, close_tag))