Testing LaminDB’s select statements#

Mock a simple LaminDB instance#

!lamin init --storage sqlapi

💬 Not registering instance on hub, if you want, call `lamin register`

💬 Loading schema modules: core==0.34.0

✅ Loaded instance: testuser1/sqlapi
✅ Created & loaded instance: testuser1/sqlapi

import lamindb as ln
import sqlmodel as sqm
import pandas as pd
from datetime import datetime
from typing import Union, Tuple, Optional, List

✅ Loaded instance: testuser1/sqlapi

class Hero(sqm.SQLModel, table=True):
    id: int = sqm.Field(primary_key=True)
    name: str
    team_id: Optional[int] = sqm.Field(foreign_key="team.id")
    team: "Team" = sqm.Relationship(back_populates="heroes")
    created_at: datetime = sqm.Field(default_factory=datetime.utcnow)

class Team(sqm.SQLModel, table=True):
    id: int = sqm.Field(primary_key=True)
    name: str
    heroes: List[Hero] = sqm.Relationship(back_populates="team")

# need to delete the instance here as we're not interested in the default LaminDB models
!rm sqlapi/sqlapi.lndb

engine = sqm.create_engine("sqlite:///sqlapi/sqlapi.lndb")

sqm.SQLModel.metadata.create_all(engine)

Add test data#

with sqm.Session(engine) as session:
    team_1 = Team(id=0, name="stars")
    session.add(team_1)
    team_2 = Team(id=1, name="moons")
    session.add(team_2)
    session.add(Hero(id=0, name="test1", team_id=team_1.id))
    session.add(Hero(id=1, name="test2", team_id=team_1.id))
    session.add(Hero(id=2, name="test3"))
    session.commit()

Test select statements#

ln.select(Hero, id=0).one()

Hero(id=0, name='test1', team_id=0, created_at=datetime.datetime(2023, 5, 30, 20, 26, 9, 816440))

ln.select(Hero).df()

	name	team_id	created_at
id
0	test1	0.0	2023-05-30 20:26:09.816440
1	test2	0.0	2023-05-30 20:26:09.816540
2	test3	NaN	2023-05-30 20:26:09.816623

ln.select(Hero, team_id=0).all()

[Hero(id=0, name='test1', team_id=0, created_at=datetime.datetime(2023, 5, 30, 20, 26, 9, 816440)),
 Hero(id=1, name='test2', team_id=0, created_at=datetime.datetime(2023, 5, 30, 20, 26, 9, 816540))]

ln.select(Hero).where(Hero.name == "test1").df()

	name	team_id	created_at
id
0	test1	0	2023-05-30 20:26:09.816440

ln.select(Hero).where(Hero.name == "test1").df()

	name	team_id	created_at
id
0	test1	0	2023-05-30 20:26:09.816440

ln.select(Hero).where(Hero.id == 0, Hero.name == "test1").df()

	name	team_id	created_at
id
0	test1	0	2023-05-30 20:26:09.816440

ln.select(Hero).where(Hero.id == 0, Hero.created_at <= datetime.utcnow()).df()

	name	team_id	created_at
id
0	test1	0	2023-05-30 20:26:09.816440

ln.select(Hero).where(sqm.or_(Hero.name == "test1", Hero.name == "test2")).df()

	name	team_id	created_at
id
0	test1	0	2023-05-30 20:26:09.816440
1	test2	0	2023-05-30 20:26:09.816540

ln.select(Hero).where(sqm.or_(Hero.name == "test1", Hero.name == "test2")).offset(
    1
).df()

	name	team_id	created_at
id
1	test2	0	2023-05-30 20:26:09.816540

ln.select(Hero).where(sqm.or_(Hero.name == "test1", Hero.name == "test2")).offset(
    0
).limit(1).df()

	name	team_id	created_at
id
0	test1	0	2023-05-30 20:26:09.816440

ln.select(Hero).where(sqm.or_(Hero.name == "test1", Hero.name == "test2")).order_by(
    sqm.desc(Hero.created_at)
).df()

	name	team_id	created_at
id
1	test2	0	2023-05-30 20:26:09.816540
0	test1	0	2023-05-30 20:26:09.816440

ln.select(Hero).join(Team).df()

	name	team_id	created_at
id
0	test1	0	2023-05-30 20:26:09.816440
1	test2	0	2023-05-30 20:26:09.816540

ln.select(Hero, Team).join(Team).all()

[(Hero(id=0, name='test1', team_id=0, created_at=datetime.datetime(2023, 5, 30, 20, 26, 9, 816440)), Team(id=0, name='stars')),
 (Hero(id=1, name='test2', team_id=0, created_at=datetime.datetime(2023, 5, 30, 20, 26, 9, 816540)), Team(id=0, name='stars'))]

dfs = ln.select(Hero, Team).join(Team).df()

for df in dfs:
    display(df)

('hero', 'id')

('hero', 'name')

('hero', 'team_id')

('hero', 'created_at')

('team', 'id')

('team', 'name')

ln.select(Hero, Team).where(Hero.team_id == Team.id).all()

[(Hero(id=0, name='test1', team_id=0, created_at=datetime.datetime(2023, 5, 30, 20, 26, 9, 816440)), Team(id=0, name='stars')),
 (Hero(id=1, name='test2', team_id=0, created_at=datetime.datetime(2023, 5, 30, 20, 26, 9, 816540)), Team(id=0, name='stars'))]

Test autoflush#

hero = Hero(name="test4")

hero

Hero(name='test4', created_at=datetime.datetime(2023, 5, 30, 20, 26, 10, 107657))

assert ln.select(Hero, name="test4").one_or_none() is None

team = ln.select(Team, name="stars").one()

hero.team = team

assert ln.select(Hero, name="test4").one_or_none() is None