feat: add select_one_record & FORCE INDEX

Author: NightMarcher

README.md

@@ -50,13 +50,14 @@ await AccountMgr.select_custom_fields(
     # wheres=Q(Q(id__in=aids), Q(gender=1), join_type="AND"),
     # wheres={"id__in": aids, "gender": 1},
     # wheres=[Q(id__in=aids), Q(gender=1)],
+    index="PRIMARY",
 )
 ```
 Generate sql and execute
 ```sql
     SELECT
       id, extend ->> '$.last_login.ipv4' ipv4, extend ->> '$.last_login.start_datetime' start_datetime, CAST(extend ->> '$.last_login.online_sec' AS SIGNED) online_sec
-    FROM account
+    FROM account FORCE INDEX (`PRIMARY`)
     WHERE id IN (1,2,3) AND gender=1
 ```
 

examples/service/routers/account.py

@@ -45,7 +45,11 @@ async def update_view(
 async def query_by_id_view(
     aid: int = Query(...),
 ):
-    account = await AccountMgr.get_by_pk(aid)
+    # account = await AccountMgr.get_by_pk(aid)
+    account = await AccountMgr.select_one_record(
+        fields=["*"],
+        wheres=Q(id=aid),
+    )
     return {"account": account}
 
 

fastapi_esql/orm/base_manager.py

@@ -63,17 +63,19 @@ async def select_custom_fields(
         cls,
         fields: List[str],
         wheres: Union[str, Q, Dict[str, Any], List[Q]],
+        index: Optional[str] = None,
         groups: Optional[List[str]] = None,
         having: Optional[str] = None,
         orders: Optional[List[str]] = None,
         offset: Optional[int] = None,
-        limit: int = 0,
+        limit: Optional[int] = None,
         conn: Optional[BaseDBAsyncClient] = None,
     ):
         sql = SQLizer.select_custom_fields(
             cls.table,
             fields,
             wheres,
+            index,
             groups,
             having,
             orders,
@@ -84,6 +86,32 @@ async def select_custom_fields(
         conn = conn or cls.ro_conn
         return await CursorHandler.fetch_dicts(sql, conn, logger)
 
+    @classmethod
+    async def select_one_record(
+        cls,
+        fields: List[str],
+        wheres: Union[str, Q, Dict[str, Any], List[Q]],
+        index: Optional[str] = None,
+        groups: Optional[List[str]] = None,
+        having: Optional[str] = None,
+        orders: Optional[List[str]] = None,
+        conn: Optional[BaseDBAsyncClient] = None,
+    ):
+        sql = SQLizer.select_custom_fields(
+            cls.table,
+            fields,
+            wheres,
+            index,
+            groups,
+            having,
+            orders,
+            0,
+            1,
+            cls.model,
+        )
+        conn = conn or cls.ro_conn
+        return await CursorHandler.fetch_one(sql, conn, logger)
+
     @classmethod
     async def update_json_field(
         cls,

fastapi_esql/utils/cursor_handler.py

@@ -19,6 +19,22 @@ async def fetch_dicts(
             logger.exception(f"{e} SQL=>{sql}")
             return None
 
+    @classmethod
+    async def fetch_one(
+        cls,
+        sql: str,
+        conn: BaseDBAsyncClient,
+        logger: Logger,
+    ) -> Optional[Dict[str, Any]]:
+        try:
+            dicts = await conn.execute_query_dict(sql)
+            if dicts:
+                return dicts[0]
+            return {}
+        except Exception as e:
+            logger.exception(f"{e} SQL=>{sql}")
+            return None
+
     @classmethod
     async def sum_row_cnt(
         cls,

fastapi_esql/utils/sqlizer.py

@@ -114,11 +114,12 @@ def select_custom_fields(
         table: str,
         fields: List[str],
         wheres: Union[str, Q, Dict[str, Any], List[Q]],
+        index: Optional[str] = None,
         groups: Optional[List[str]] = None,
         having: Optional[str] = None,
         orders: Optional[List[str]] = None,
         offset: Optional[int] = None,
-        limit: int = 0,
+        limit: Optional[int] = None,
         model: Optional[Model] = None,
     ) -> Optional[str]:
         if not all([table, fields, wheres]):
@@ -136,11 +137,12 @@ def select_custom_fields(
         sql = """
     SELECT
       {}
-    FROM {}
+    FROM {}{}
     WHERE {}
 {}""".format(
         ", ".join(fields),
         table,
+        f" FORCE INDEX (`{index}`)" if index else "",
         cls.resolve_wheres(wheres, model),
         "\n".join(i for i in extras if i),
     )

tests/test_cursor_handler.py

@@ -21,6 +21,18 @@ async def test_fetch_dicts(self):
         print(f"result => {result}")
         assert result
 
+    async def test_fetch_one(self):
+        one = await CursorHandler.fetch_one("SELECT 1 idx", self.conn, logger)
+        assert one == {"idx": 1}
+
+        with patch(
+            "tortoise.backends.mysql.client_class.execute_query_dict",
+        ) as mock_exec:
+            mock_exec.return_value = []
+            empty = await CursorHandler.fetch_one("SELECT 1 idx", self.conn, logger)
+            print(f"empty => {empty}")
+            assert empty == {}
+
     @patch("tortoise.backends.mysql.client_class.execute_query")
     async def test_sum_row_cnt(self, mock_exec):
         mock_exec.return_value = 10, object()

tests/test_sqlizer.py

@@ -114,12 +114,13 @@ def test_select_custom_fields(self):
                 "CAST(extend ->> '$.last_login.online_sec' AS SIGNED) online_sec"
             ],
             wheres=f"id IN ({','.join(map(str, aids))}) AND gender=1",
+            index="PRIMARY",
             model=self.model,
         )
         assert basic_sql == """
     SELECT
       id, extend ->> '$.last_login.ipv4' ipv4, extend ->> '$.last_login.start_datetime' start_datetime, CAST(extend ->> '$.last_login.online_sec' AS SIGNED) online_sec
-    FROM account
+    FROM account FORCE INDEX (`PRIMARY`)
     WHERE id IN (1,2,3) AND gender=1
 """