From 7c3e2a67bd5fe0a92d9862a6ba95ead3cdf5e9c8 Mon Sep 17 00:00:00 2001 From: long2ice Date: Thu, 2 Sep 2021 10:30:18 +0800 Subject: [PATCH] Add `Model.raw` method to support the raw sql query --- CHANGELOG.rst | 5 +++++ pyproject.toml | 2 +- tests/test_model_methods.py | 7 ++++++ tortoise/__init__.py | 2 +- tortoise/backends/base/executor.py | 6 ++++-- tortoise/models.py | 15 ++++++++++++- tortoise/queryset.py | 34 +++++++++++++++++++++++++++++- 7 files changed, 65 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7d81cd78a..2e237537e 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -9,6 +9,11 @@ Changelog 0.17 ==== + +0.17.8 +------ +- Add `Model.raw` method to support the raw sql query. + 0.17.7 ------ - Fix `select_related` behaviour for forward relation. (#825) diff --git a/pyproject.toml b/pyproject.toml index c346ee449..215129b2b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "tortoise-orm" -version = "0.17.7" +version = "0.17.8" description = "Easy async ORM for python, built with relations in mind" authors = ["Andrey Bondar ", "Nickolas Grigoriadis ", "long2ice "] license = "Apache-2.0" diff --git a/tests/test_model_methods.py b/tests/test_model_methods.py index 223bc9bfb..7251d0c21 100644 --- a/tests/test_model_methods.py +++ b/tests/test_model_methods.py @@ -288,6 +288,13 @@ async def test_force_update_raise(self): with self.assertRaises(IntegrityError): await obj.save(force_update=True) + async def test_raw(self): + await self.cls.create(name="TestRaw", id=self.mdl.id) + ret = await self.cls.raw("select * from tournament where name='TestRaw'") + self.assertEqual(len(ret), 1) + ret = await self.cls.raw("select * from tournament where name='111'") + self.assertEqual(len(ret), 0) + class TestModelMethodsNoID(TestModelMethods): async def setUp(self): diff --git a/tortoise/__init__.py b/tortoise/__init__.py index b1ca97365..0ef6047af 100644 --- a/tortoise/__init__.py +++ b/tortoise/__init__.py @@ -709,4 +709,4 @@ async def do_stuff(): loop.run_until_complete(Tortoise.close_connections()) -__version__ = "0.17.7" +__version__ = "0.17.8" diff --git a/tortoise/backends/base/executor.py b/tortoise/backends/base/executor.py index 431771ec8..9fafb0210 100644 --- a/tortoise/backends/base/executor.py +++ b/tortoise/backends/base/executor.py @@ -23,7 +23,7 @@ from pypika.terms import ArithmeticExpression, Function from tortoise.exceptions import OperationalError -from tortoise.expressions import F +from tortoise.expressions import F, RawSQL from tortoise.fields.base import Field from tortoise.fields.relational import ( BackwardFKRelation, @@ -122,7 +122,9 @@ async def execute_explain(self, query: Query) -> Any: sql = " ".join((self.EXPLAIN_PREFIX, query.get_sql())) return (await self.db.execute_query(sql))[1] - async def execute_select(self, query: Query, custom_fields: Optional[list] = None) -> list: + async def execute_select( + self, query: Union[Query, RawSQL], custom_fields: Optional[list] = None + ) -> list: _, raw_results = await self.db.execute_query(query.get_sql()) instance_list = [] for row in raw_results: diff --git a/tortoise/models.py b/tortoise/models.py index e906ec4b9..f83bd67ed 100644 --- a/tortoise/models.py +++ b/tortoise/models.py @@ -48,7 +48,7 @@ from tortoise.functions import Function from tortoise.indexes import Index from tortoise.manager import Manager -from tortoise.queryset import ExistsQuery, Q, QuerySet, QuerySetSingle +from tortoise.queryset import ExistsQuery, Q, QuerySet, QuerySetSingle, RawSQLQuery from tortoise.router import router from tortoise.signals import Signals from tortoise.transactions import current_transaction_map, in_transaction @@ -1203,6 +1203,19 @@ def get(cls: Type[MODEL], *args: Q, **kwargs: Any) -> QuerySetSingle[MODEL]: """ return cls._meta.manager.get_queryset().get(*args, **kwargs) + @classmethod + def raw(cls, sql: str) -> "RawSQLQuery": + """ + Executes a RAW SQL and returns the result + + .. code-block:: python3 + + result = await User.raw("select * from users where name like '%test%'") + + :param sql: The raw sql. + """ + return cls._meta.manager.get_queryset().raw(sql) + @classmethod def exists(cls: Type[MODEL], *args: Q, **kwargs: Any) -> ExistsQuery: """ diff --git a/tortoise/queryset.py b/tortoise/queryset.py index 60244930a..473be0742 100644 --- a/tortoise/queryset.py +++ b/tortoise/queryset.py @@ -33,7 +33,7 @@ MultipleObjectsReturned, ParamsError, ) -from tortoise.expressions import F +from tortoise.expressions import F, RawSQL from tortoise.fields.relational import ( ForeignKeyFieldInstance, OneToOneFieldInstance, @@ -649,6 +649,12 @@ def all(self) -> "QuerySet[MODEL]": """ return self._clone() + def raw(self, sql: str) -> "RawSQLQuery": + """ + Return the QuerySet from raw SQL + """ + return RawSQLQuery(model=self.model, db=self._db, sql=sql) + def first(self) -> QuerySetSingle[Optional[MODEL]]: """ Limit queryset to one object and return one object instead of list. @@ -1497,3 +1503,29 @@ async def _execute(self) -> List[dict]: row[col] = func(row[col]) return result + + +class RawSQLQuery(AwaitableQuery): + + __slots__ = ("_sql", "_db") + + def __init__(self, model: Type[MODEL], db: BaseDBAsyncClient, sql: str): + super().__init__(model) + self._sql = sql + self._db = db + + def _make_query(self) -> None: + self.query = RawSQL(self._sql) + + async def _execute(self) -> Any: + instance_list = await self._db.executor_class( + model=self.model, + db=self._db, + ).execute_select(self.query) + return instance_list + + def __await__(self) -> Generator[Any, None, List[MODEL]]: + if self._db is None: + self._db = self._choose_db() # type: ignore + self._make_query() + return self._execute().__await__()