Add find_one() implementation
This commit is contained in:
		
							parent
							
								
									f8c55236c1
								
							
						
					
					
						commit
						85ba111260
					
				
					 2 changed files with 33 additions and 38 deletions
				
			
		| 
						 | 
					@ -85,6 +85,9 @@ class Expression:
 | 
				
			||||||
    def __and__(self, other):
 | 
					    def __and__(self, other):
 | 
				
			||||||
        return Expression(left=self, op=Operators.AND, right=other)
 | 
					        return Expression(left=self, op=Operators.AND, right=other)
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					    def __or__(self, other):
 | 
				
			||||||
 | 
					        return Expression(left=self, op=Operators.OR, right=other)
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					
 | 
				
			||||||
ExpressionOrNegated = Union[Expression, NegatedExpression]
 | 
					ExpressionOrNegated = Union[Expression, NegatedExpression]
 | 
				
			||||||
 | 
					
 | 
				
			||||||
| 
						 | 
					@ -109,11 +112,15 @@ class FindQuery:
 | 
				
			||||||
    expressions: Sequence[Expression]
 | 
					    expressions: Sequence[Expression]
 | 
				
			||||||
    expression: Expression = dataclasses.field(init=False)
 | 
					    expression: Expression = dataclasses.field(init=False)
 | 
				
			||||||
    query: str = dataclasses.field(init=False)
 | 
					    query: str = dataclasses.field(init=False)
 | 
				
			||||||
 | 
					    pagination: List[str] = dataclasses.field(init=False)
 | 
				
			||||||
    model: Type['RedisModel']
 | 
					    model: Type['RedisModel']
 | 
				
			||||||
 | 
					    limit: Optional[int] = None
 | 
				
			||||||
 | 
					    offset: Optional[int] = None
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    def __post_init__(self):
 | 
					    def __post_init__(self):
 | 
				
			||||||
        self.expression = reduce(operator.and_, self.expressions)
 | 
					        self.expression = reduce(operator.and_, self.expressions)
 | 
				
			||||||
        self.query = self.resolve_redisearch_query(self.expression)
 | 
					        self.query = self.resolve_redisearch_query(self.expression)
 | 
				
			||||||
 | 
					        self.pagination = self.resolve_redisearch_pagination()
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    def resolve_field_type(self, field: ModelField) -> RediSearchFieldTypes:
 | 
					    def resolve_field_type(self, field: ModelField) -> RediSearchFieldTypes:
 | 
				
			||||||
        if getattr(field.field_info, 'primary_key', None):
 | 
					        if getattr(field.field_info, 'primary_key', None):
 | 
				
			||||||
| 
						 | 
					@ -159,6 +166,14 @@ class FindQuery:
 | 
				
			||||||
 | 
					
 | 
				
			||||||
        return result
 | 
					        return result
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					    def resolve_redisearch_pagination(self):
 | 
				
			||||||
 | 
					        """Resolve pagination options for a query."""
 | 
				
			||||||
 | 
					        if not self.limit and not self.offset:
 | 
				
			||||||
 | 
					            return []
 | 
				
			||||||
 | 
					        offset = self.offset or 0
 | 
				
			||||||
 | 
					        limit = self.limit or 10
 | 
				
			||||||
 | 
					        return ["LIMIT", offset, limit]
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    def resolve_redisearch_query(self, expression: ExpressionOrNegated):
 | 
					    def resolve_redisearch_query(self, expression: ExpressionOrNegated):
 | 
				
			||||||
        """Resolve an expression to a string RediSearch query."""
 | 
					        """Resolve an expression to a string RediSearch query."""
 | 
				
			||||||
        field_type = None
 | 
					        field_type = None
 | 
				
			||||||
| 
						 | 
					@ -210,8 +225,11 @@ class FindQuery:
 | 
				
			||||||
        return result
 | 
					        return result
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    def find(self):
 | 
					    def find(self):
 | 
				
			||||||
        return self.model.db().execute_command("ft.search", self.model.Meta.index_name,
 | 
					        args = ["ft.search", self.model.Meta.index_name, self.query]
 | 
				
			||||||
                                               self.query)
 | 
					        # TODO: Do we need self.pagination if we're just appending to query anyway?
 | 
				
			||||||
 | 
					        if self.pagination:
 | 
				
			||||||
 | 
					            args.extend(self.pagination)
 | 
				
			||||||
 | 
					        return self.model.db().execute_command(*args)
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					
 | 
				
			||||||
class PrimaryKeyCreator(Protocol):
 | 
					class PrimaryKeyCreator(Protocol):
 | 
				
			||||||
| 
						 | 
					@ -518,8 +536,10 @@ class RedisModel(BaseModel, abc.ABC, metaclass=ModelMeta):
 | 
				
			||||||
        return cls.from_redis(raw_result)
 | 
					        return cls.from_redis(raw_result)
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    @classmethod
 | 
					    @classmethod
 | 
				
			||||||
    def find_one(cls, *expressions: Sequence[Expression]):
 | 
					    def find_one(cls, *expressions: Expression):
 | 
				
			||||||
        return cls
 | 
					        query = FindQuery(expressions=expressions, model=cls, limit=1, offset=0)
 | 
				
			||||||
 | 
					        raw_result = query.find()
 | 
				
			||||||
 | 
					        return cls.from_redis(raw_result)[0]
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    @classmethod
 | 
					    @classmethod
 | 
				
			||||||
    def add(cls, models: Sequence['RedisModel']) -> Sequence['RedisModel']:
 | 
					    def add(cls, models: Sequence['RedisModel']) -> Sequence['RedisModel']:
 | 
				
			||||||
| 
						 | 
					
 | 
				
			||||||
| 
						 | 
					@ -162,8 +162,6 @@ def test_exact_match_queries():
 | 
				
			||||||
    member2.save()
 | 
					    member2.save()
 | 
				
			||||||
    member3.save()
 | 
					    member3.save()
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    import ipdb; ipdb.set_trace()
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
    # # TODO: How to help IDEs know that last_name is not a str, but a wrapped expression?
 | 
					    # # TODO: How to help IDEs know that last_name is not a str, but a wrapped expression?
 | 
				
			||||||
    actual = Member.find(Member.last_name == "Brookins")
 | 
					    actual = Member.find(Member.last_name == "Brookins")
 | 
				
			||||||
    assert actual == [member2, member1]
 | 
					    assert actual == [member2, member1]
 | 
				
			||||||
| 
						 | 
					@ -179,15 +177,14 @@ def test_exact_match_queries():
 | 
				
			||||||
    actual = Member.find(Member.last_name != "Brookins")
 | 
					    actual = Member.find(Member.last_name != "Brookins")
 | 
				
			||||||
    assert actual == [member3]
 | 
					    assert actual == [member3]
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					    actual = Member.find(
 | 
				
			||||||
 | 
					        (Member.last_name == "Brookins") & (Member.first_name == "Andrew")
 | 
				
			||||||
 | 
					        | (Member.first_name == "Kim")
 | 
				
			||||||
 | 
					    )
 | 
				
			||||||
 | 
					    assert actual == [member2, member1]
 | 
				
			||||||
 | 
					
 | 
				
			||||||
# actual = Member.find(
 | 
					    actual = Member.find_one(Member.last_name == "Brookins")
 | 
				
			||||||
    #     (Member.last_name == "Brookins") & (Member.first_name == "Andrew")
 | 
					    assert actual == member2
 | 
				
			||||||
    #     | (Member.first_name == "Kim")
 | 
					 | 
				
			||||||
    # )
 | 
					 | 
				
			||||||
    # assert actual == [member1, member2]
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
    # actual = Member.find_one(Member.last_name == "Brookins")
 | 
					 | 
				
			||||||
    # assert actual == member1
 | 
					 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					
 | 
				
			||||||
def test_schema():
 | 
					def test_schema():
 | 
				
			||||||
| 
						 | 
					@ -196,28 +193,6 @@ def test_schema():
 | 
				
			||||||
        an_integer: int
 | 
					        an_integer: int
 | 
				
			||||||
        a_float: float
 | 
					        a_float: float
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    assert Address.schema() == "SCHEMA pk TAG SORTABLE a_string TEXT an_integer NUMERIC " \
 | 
					    # TODO: Fix
 | 
				
			||||||
 | 
					    assert Address.schema() == "ON HASH PREFIX 1 redis-developer:basehashmodel: SCHEMA pk TAG SORTABLE a_string TEXT an_integer NUMERIC " \
 | 
				
			||||||
                               "a_float NUMERIC"
 | 
					                               "a_float NUMERIC"
 | 
				
			||||||
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
# ---
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
from typing import Optional
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
from sqlmodel import Field, Session, SQLModel, create_engine, select
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
class Hero(SQLModel, table=True):
 | 
					 | 
				
			||||||
    id: Optional[int] = Field(default=None, primary_key=True)
 | 
					 | 
				
			||||||
    name: str
 | 
					 | 
				
			||||||
    secret_name: str
 | 
					 | 
				
			||||||
    age: Optional[int] = None
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
engine = create_engine("sqlite:///database.db")
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
with Session(engine) as session:
 | 
					 | 
				
			||||||
    import ipdb; ipdb.set_trace()
 | 
					 | 
				
			||||||
    statement = select(Hero).where(Hero.name == "Spider-Boy")
 | 
					 | 
				
			||||||
    hero = session.exec(statement).first()
 | 
					 | 
				
			||||||
    print(hero)
 | 
					 | 
				
			||||||
| 
						 | 
					
 | 
				
			||||||
		Loading…
	
	Add table
		Add a link
		
	
		Reference in a new issue