[DOP-28871] Improve SQLAlchemy query caching

Author: dolfinus

data_rentgen/db/repositories/dataset.py

@@ -13,6 +13,7 @@
     String,
     any_,
     asc,
+    bindparam,
     cast,
     desc,
     distinct,
@@ -28,25 +29,55 @@
 from data_rentgen.db.utils.search import make_tsquery, ts_match, ts_rank
 from data_rentgen.dto import DatasetDTO, PaginationDTO
 
+fetch_bulk_query = select(Dataset).where(
+    tuple_(Dataset.location_id, func.lower(Dataset.name)).in_(
+        select(
+            func.unnest(
+                cast(bindparam("location_ids"), ARRAY(Integer())),
+                cast(bindparam("names"), ARRAY(String())),
+            )
+            .table_valued("location_id", "name")
+            .render_derived(),
+        ),
+    ),
+)
+
+get_list_query = (
+    select(Dataset)
+    .where(Dataset.id == any_(bindparam("dataset_ids")))
+    .options(selectinload(Dataset.location).selectinload(Location.addresses))
+    .options(selectinload(Dataset.tag_values).selectinload(TagValue.tag))
+)
+
+get_one_query = select(Dataset).where(
+    Dataset.location_id == bindparam("location_id"),
+    func.lower(Dataset.name) == bindparam("name"),
+)
+
+get_stats_query = (
+    select(
+        Dataset.location_id.label("location_id"),
+        func.count(Dataset.id.distinct()).label("total_datasets"),
+    )
+    .where(
+        Dataset.location_id == any_(bindparam("location_ids")),
+    )
+    .group_by(Dataset.location_id)
+)
+
 
 class DatasetRepository(Repository[Dataset]):
     async def fetch_bulk(self, datasets_dto: list[DatasetDTO]) -> list[tuple[DatasetDTO, Dataset | None]]:
         if not datasets_dto:
             return []
 
-        location_ids = [dataset_dto.location.id for dataset_dto in datasets_dto]
-        names = [dataset_dto.name.lower() for dataset_dto in datasets_dto]
-        pairs = (
-            func.unnest(
-                cast(location_ids, ARRAY(Integer())),
-                cast(names, ARRAY(String())),
-            )
-            .table_valued("location_id", "name")
-            .render_derived()
+        scalars = await self._session.scalars(
+            fetch_bulk_query,
+            {
+                "location_ids": [item.location.id for item in datasets_dto],
+                "names": [item.name.lower() for item in datasets_dto],
+            },
         )
-
-        statement = select(Dataset).where(tuple_(Dataset.location_id, func.lower(Dataset.name)).in_(select(pairs)))
-        scalars = await self._session.scalars(statement)
         existing = {(dataset.location_id, dataset.name.lower()): dataset for dataset in scalars.all()}
         return [
             (
@@ -139,39 +170,24 @@ async def paginate(
     async def list_by_ids(self, dataset_ids: Collection[int]) -> list[Dataset]:
         if not dataset_ids:
             return []
-        query = (
-            select(Dataset)
-            .where(Dataset.id == any_(list(dataset_ids)))  # type: ignore[arg-type]
-            .options(selectinload(Dataset.location).selectinload(Location.addresses))
-            .options(selectinload(Dataset.tag_values).selectinload(TagValue.tag))
-        )
-        result = await self._session.scalars(query)
+        result = await self._session.scalars(get_list_query, {"dataset_ids": list(dataset_ids)})
         return list(result.all())
 
     async def get_stats_by_location_ids(self, location_ids: Collection[int]) -> dict[int, Row]:
         if not location_ids:
             return {}
 
-        query = (
-            select(
-                Dataset.location_id.label("location_id"),
-                func.count(Dataset.id.distinct()).label("total_datasets"),
-            )
-            .where(
-                Dataset.location_id == any_(list(location_ids)),  # type: ignore[arg-type]
-            )
-            .group_by(Dataset.location_id)
-        )
-
-        query_result = await self._session.execute(query)
+        query_result = await self._session.execute(get_stats_query, {"location_ids": list(location_ids)})
         return {row.location_id: row for row in query_result.all()}
 
     async def _get(self, dataset: DatasetDTO) -> Dataset | None:
-        statement = select(Dataset).where(
-            Dataset.location_id == dataset.location.id,
-            func.lower(Dataset.name) == dataset.name.lower(),
+        return await self._session.scalar(
+            get_one_query,
+            {
+                "location_id": dataset.location.id,
+                "name": dataset.name.lower(),
+            },
         )
-        return await self._session.scalar(statement)
 
     async def _create(self, dataset: DatasetDTO) -> Dataset:
         result = Dataset(location_id=dataset.location.id, name=dataset.name)

data_rentgen/db/repositories/dataset_symlink.py

@@ -3,12 +3,37 @@
 
 from collections.abc import Collection
 
-from sqlalchemy import ARRAY, BindParameter, Integer, any_, bindparam, cast, func, or_, select, tuple_
+from sqlalchemy import ARRAY, Integer, any_, bindparam, cast, func, or_, select, tuple_
 
 from data_rentgen.db.models.dataset_symlink import DatasetSymlink, DatasetSymlinkType
 from data_rentgen.db.repositories.base import Repository
 from data_rentgen.dto import DatasetSymlinkDTO
 
+fetch_bulk_query = select(DatasetSymlink).where(
+    tuple_(DatasetSymlink.from_dataset_id, DatasetSymlink.to_dataset_id).in_(
+        select(
+            func.unnest(
+                cast(bindparam("from_dataset_ids"), ARRAY(Integer())),
+                cast(bindparam("to_dataset_ids"), ARRAY(Integer())),
+            )
+            .table_valued("from_dataset_ids", "to_dataset_ids")
+            .render_derived(),
+        ),
+    ),
+)
+
+get_list_query = select(DatasetSymlink).where(
+    or_(
+        DatasetSymlink.from_dataset_id == any_(bindparam("dataset_ids")),
+        DatasetSymlink.to_dataset_id == any_(bindparam("dataset_ids")),
+    ),
+)
+
+get_one_query = select(DatasetSymlink).where(
+    DatasetSymlink.from_dataset_id == bindparam("from_dataset_id"),
+    DatasetSymlink.to_dataset_id == bindparam("to_dataset_id"),
+)
+
 
 class DatasetSymlinkRepository(Repository[DatasetSymlink]):
     async def fetch_bulk(
@@ -18,22 +43,13 @@ async def fetch_bulk(
         if not dataset_symlinks_dto:
             return []
 
-        from_dataset_ids = [dataset_symlink_dto.from_dataset.id for dataset_symlink_dto in dataset_symlinks_dto]
-        to_dataset_ids = [dataset_symlink_dto.to_dataset.id for dataset_symlink_dto in dataset_symlinks_dto]
-
-        pairs = (
-            func.unnest(
-                cast(from_dataset_ids, ARRAY(Integer())),
-                cast(to_dataset_ids, ARRAY(Integer())),
-            )
-            .table_valued("from_dataset_ids", "to_dataset_ids")
-            .render_derived()
+        scalars = await self._session.scalars(
+            fetch_bulk_query,
+            {
+                "from_dataset_ids": [item.from_dataset.id for item in dataset_symlinks_dto],
+                "to_dataset_ids": [item.to_dataset.id for item in dataset_symlinks_dto],
+            },
         )
-
-        statement = select(DatasetSymlink).where(
-            tuple_(DatasetSymlink.from_dataset_id, DatasetSymlink.to_dataset_id).in_(select(pairs)),
-        )
-        scalars = await self._session.scalars(statement)
         existing = {(item.from_dataset_id, item.to_dataset_id): item for item in scalars.all()}
         return [
             (
@@ -52,22 +68,17 @@ async def list_by_dataset_ids(self, dataset_ids: Collection[int]) -> list[Datase
         if not dataset_ids:
             return []
 
-        param: BindParameter[list[int]] = bindparam("dataset_ids")
-        query = select(DatasetSymlink).where(
-            or_(
-                DatasetSymlink.from_dataset_id == any_(param),
-                DatasetSymlink.to_dataset_id == any_(param),
-            ),
-        )
-        scalars = await self._session.scalars(query, {"dataset_ids": list(dataset_ids)})
+        scalars = await self._session.scalars(get_list_query, {"dataset_ids": list(dataset_ids)})
         return list(scalars.all())
 
     async def _get(self, dataset_symlink: DatasetSymlinkDTO) -> DatasetSymlink | None:
-        query = select(DatasetSymlink).where(
-            DatasetSymlink.from_dataset_id == dataset_symlink.from_dataset.id,
-            DatasetSymlink.to_dataset_id == dataset_symlink.to_dataset.id,
+        return await self._session.scalar(
+            get_one_query,
+            {
+                "from_dataset_id": dataset_symlink.from_dataset.id,
+                "to_dataset_id": dataset_symlink.to_dataset.id,
+            },
         )
-        return await self._session.scalar(query)
 
     async def _create(self, dataset_symlink: DatasetSymlinkDTO) -> DatasetSymlink:
         result = DatasetSymlink(

data_rentgen/db/repositories/input.py

@@ -16,6 +16,17 @@
     extract_timestamp_from_uuid,
 )
 
+insert_statement = insert(Input)
+inserted_row = insert_statement.excluded
+insert_statement = insert_statement.on_conflict_do_update(
+    index_elements=[Input.created_at, Input.id],
+    set_={
+        "num_bytes": func.greatest(inserted_row.num_bytes, Input.num_bytes),
+        "num_rows": func.greatest(inserted_row.num_rows, Input.num_rows),
+        "num_files": func.greatest(inserted_row.num_files, Input.num_files),
+    },
+)
+
 
 @dataclass
 class InputRow:
@@ -37,19 +48,8 @@ async def create_or_update_bulk(self, inputs: list[InputDTO]) -> None:
         if not inputs:
             return
 
-        insert_statement = insert(Input)
-        new_row = insert_statement.excluded
-        statement = insert_statement.on_conflict_do_update(
-            index_elements=[Input.created_at, Input.id],
-            set_={
-                "num_bytes": func.greatest(new_row.num_bytes, Input.num_bytes),
-                "num_rows": func.greatest(new_row.num_rows, Input.num_rows),
-                "num_files": func.greatest(new_row.num_files, Input.num_files),
-            },
-        )
-
         await self._session.execute(
-            statement,
+            insert_statement,
             [
                 {
                     "id": item.generate_id(),

data_rentgen/db/repositories/job.py

@@ -13,6 +13,7 @@
     String,
     any_,
     asc,
+    bindparam,
     cast,
     desc,
     func,
@@ -30,6 +31,44 @@
 UNKNOWN_JOB_TYPE = 0
 
 
+fetch_bulk_query = select(Job).where(
+    tuple_(Job.location_id, func.lower(Job.name)).in_(
+        select(
+            func.unnest(
+                cast(bindparam("location_ids"), ARRAY(Integer())),
+                cast(bindparam("names"), ARRAY(String())),
+            )
+            .table_valued("location_id", "name")
+            .render_derived(),
+        ),
+    ),
+)
+
+get_one_query = select(Job).where(
+    Job.location_id == bindparam("location_id"),
+    func.lower(Job.name) == bindparam("name"),
+)
+
+get_list_query = (
+    select(Job)
+    .where(
+        Job.id == any_(bindparam("job_ids")),
+    )
+    .options(selectinload(Job.location).selectinload(Location.addresses))
+)
+
+get_stats_query = (
+    select(
+        Job.location_id.label("location_id"),
+        func.count(Job.id.distinct()).label("total_jobs"),
+    )
+    .where(
+        Job.location_id == any_(bindparam("location_ids")),
+    )
+    .group_by(Job.location_id)
+)
+
+
 class JobRepository(Repository[Job]):
     async def paginate(
         self,
@@ -90,19 +129,13 @@ async def fetch_bulk(self, jobs_dto: list[JobDTO]) -> list[tuple[JobDTO, Job | N
         if not jobs_dto:
             return []
 
-        location_ids = [job_dto.location.id for job_dto in jobs_dto]
-        names = [job_dto.name.lower() for job_dto in jobs_dto]
-        pairs = (
-            func.unnest(
-                cast(location_ids, ARRAY(Integer())),
-                cast(names, ARRAY(String())),
-            )
-            .table_valued("location_id", "name")
-            .render_derived()
+        scalars = await self._session.scalars(
+            fetch_bulk_query,
+            {
+                "location_ids": [item.location.id for item in jobs_dto],
+                "names": [item.name.lower() for item in jobs_dto],
+            },
         )
-
-        statement = select(Job).where(tuple_(Job.location_id, func.lower(Job.name)).in_(select(pairs)))
-        scalars = await self._session.scalars(statement)
         existing = {(job.location_id, job.name.lower()): job for job in scalars.all()}
         return [
             (
@@ -121,11 +154,13 @@ async def create_or_update(self, job: JobDTO) -> Job:
         return await self.update(result, job)
 
     async def _get(self, job: JobDTO) -> Job | None:
-        statement = select(Job).where(
-            Job.location_id == job.location.id,
-            func.lower(Job.name) == job.name.lower(),
+        return await self._session.scalar(
+            get_one_query,
+            {
+                "location_id": job.location.id,
+                "name": job.name.lower(),
+            },
         )
-        return await self._session.scalar(statement)
 
     async def _create(self, job: JobDTO) -> Job:
         result = Job(
@@ -147,28 +182,13 @@ async def update(self, existing: Job, new: JobDTO) -> Job:
     async def list_by_ids(self, job_ids: Collection[int]) -> list[Job]:
         if not job_ids:
             return []
-        query = (
-            select(Job)
-            .where(Job.id == any_(list(job_ids)))  # type: ignore[arg-type]
-            .options(selectinload(Job.location).selectinload(Location.addresses))
-        )
-        result = await self._session.scalars(query)
+
+        result = await self._session.scalars(get_list_query, {"job_ids": list(job_ids)})
         return list(result.all())
 
     async def get_stats_by_location_ids(self, location_ids: Collection[int]) -> dict[int, Row]:
         if not location_ids:
             return {}
 
-        query = (
-            select(
-                Job.location_id.label("location_id"),
-                func.count(Job.id.distinct()).label("total_jobs"),
-            )
-            .where(
-                Job.location_id == any_(list(location_ids)),  # type: ignore[arg-type]
-            )
-            .group_by(Job.location_id)
-        )
-
-        query_result = await self._session.execute(query)
+        query_result = await self._session.execute(get_stats_query, {"location_ids": list(location_ids)})
         return {row.location_id: row for row in query_result.all()}