fastapi / sqlmodel

SQL databases in Python, designed for simplicity, compatibility, and robustness.
https://sqlmodel.tiangolo.com/
MIT License
14.57k stars 664 forks source link

How to write statements like "select count(*) from table_name where col_name>0"? #494

Open rxy1212 opened 2 years ago

rxy1212 commented 2 years ago

First Check

Commit to Help

Example Code

How to write statements like "select count(*) from table_name where col_name>0"?

Description

How to write statements like "select count(*) from table_name where col_name>0"?

Operating System

Linux

Operating System Details

No response

SQLModel Version

0.0.8

Python Version

3.7.4

Additional Context

No response

meirdev commented 2 years ago

Two options:

from typing import Optional

from sqlalchemy import func
from sqlmodel import Field, SQLModel, Session, create_engine

class Table(SQLModel, table=True):
    id: Optional[int] = Field(primary_key=True)
    column: int

engine = create_engine(
    "sqlite:///",
    echo=True,
)

with Session(engine) as session:
    SQLModel.metadata.create_all(engine)

    # Option #1:

    print(session.query(Table).where(Table.column > 10).count())

    """
    SELECT count(*) AS count_1 
    FROM (SELECT "table".id AS table_id, "table"."column" AS table_column 
    FROM "table" 
    WHERE "table"."column" > ?) AS anon_1
    """

    # Option #2:

    print(session.query(Table).with_entities(func.count()).where(Table.column > 10).scalar())

    """
    SELECT count(*) AS count_1 
    FROM "table" 
    WHERE "table"."column" > ?
    """
ThirVondukr commented 2 years ago

@rxy1212 You can use count function via sqlalchemy.func object like in regular sql:

from typing import Optional

from sqlalchemy import func
from sqlmodel import Field, SQLModel, Session, create_engine, select

class Table(SQLModel, table=True):
    id: Optional[int] = Field(primary_key=True)
    column: int

engine = create_engine(
    "sqlite:///",
    echo=True,
)

SQLModel.metadata.create_all(engine)

stmt = select(func.count(Table.id)).where(Table.column > 10)
print(stmt)
"""
SELECT count("table".id) AS count_1 
FROM "table" 
WHERE "table"."column" > :column_1
"""

with Session(engine) as session:
    for i in range(150):
        session.add(Table(column=i))
    session.flush()
    print(session.scalar(stmt)) # 139
rxy1212 commented 2 years ago

That's make sense, I'll have a try on it. Thank you @meirdev @ThirVondukr.

Metal-joker commented 1 year ago

@meirdev @ThirVondukr just out of curiosity, how you guys found this solution? It's not in the official document

antont commented 1 year ago

@Metal-joker what do you mean? I think all those are in sqlalchemy docs. SQLModel is a way to use SA with Pydantic models. https://docs.sqlalchemy.org/en/14/core/functions.html#sqlalchemy.sql.functions.count

makisukurisu commented 8 months ago

This solution is not applicable anymore. Or I'm doing something wrong, but I don't have a where method, and I get a warning:

...sqlmodel/sql/expression.py:225: SAWarning: implicitly coercing SELECT object to scalar subquery; please use the .scalar_subquery() method to produce a scalar subquery.
  return super().where(*whereclause)  # type: ignore[arg-type]

I've found this gist, and adapted it for sqlmodel, here's what I've got:

from sqlmodel import Session
from sqlmodel.sql.expression import SelectOfScalar

def get_count(session: Session, q: SelectOfScalar) -> int:
    count_q = q.with_only_columns(func.count()).order_by(None).select_from(*q.froms)
    iterator = session.exec(count_q)
    for count in iterator:
        return count
    return 0

This is a utility method, that can be used wherever we need it. For example:

import pydantic

class SearchSchema(pydantic.BaseModel):
    opt_param: str | None = None

class AccountModel(SQLModel, table=True):
    # Doesn't really matter TBH
    opt_param: str

def get_accounts(session: Session, search_params: SearchSchema)
    query = select(AccountModel)
    if search_params.opt_param:
        query = query.filter(AccountModel.opt_param == search.opt_param)

    # You can add more filters here

    total = get_count(db, query)  # Here's your result
    query = db.exec(query)
    # Rest of the method here

Works fine for me, at least - for now.

ccppoo commented 7 months ago

a bit more update from makisukurisu's comment to not get deprecate warnings

from sqlmodel.sql.expression import SelectOfScalar
from sqlmodel import func

def get_count(session: Session, q: SelectOfScalar) -> int:

    count_q = q.with_only_columns(func.count()).order_by(None).select_from(q.get_final_froms()[0])
    iterator = session.exec(count_q)
    for count in iterator:
        return count
    return 0

and use it like

# change to your code from below
with Session(engine) as sess:

    query = select(YourTable).\
                      where(YourColumn == 1)
    total = get_count(sess, query)

    print(f'{total=}')
    sess.close()
antoine-lizee commented 6 months ago

I have a simpler problem, just trying to do session.query(MyObject).count() and I get a warning:

<ipython-input-9-9445795a22fe>:1: DeprecationWarning: 
        🚨 You probably want to use `session.exec()` instead of `session.query()`.
        `session.exec()` is SQLModel's own short version with increased type
        annotations.
        Or otherwise you might want to use `session.execute()` instead of
        `session.query()`.

  session.query(MyObject).count()

What is the expressive, short way to do something similar in SQLModel?

In the spirit of using SQLModel as a thin wrapper, I'd rather continue to use the SQLA syntax than build and maintain a (rather complicated) get_count helper as above. But maybe I'm missing something?

jackbravo commented 3 months ago

SQLModel has a func.count function that is working ok:

from sqlmodel import Session, col, func, select
from .database import engine
from .models import Page

with Session(engine) as session:
    count = session.exec(select(func.count(col(Page.id))).where(Page.id > 100)).one()
Graeme22 commented 1 day ago

Here's an example of how to get both the count and the items using a single query: https://gist.github.com/Graeme22/dd5e880263f9bd2b8438cc9d079f8d38