Merge pull request #297 from eclecticiq/get-objects-performance-optimization

Get objects performance optimization

Author: eric-eclecticiq

opentaxii/persistence/sqldb/api.py

@@ -6,9 +6,10 @@
 from typing import Dict, List, Optional, Tuple, no_type_check
 
 import six
+import sqlalchemy
 import structlog
 from sqlalchemy import and_, func, literal, or_
-from sqlalchemy.orm import Query, load_only
+from sqlalchemy.orm import Query, aliased, load_only
 
 from opentaxii.common.sqldb import BaseSQLDatabaseAPI
 from opentaxii.persistence import OpenTAXII2PersistenceAPI, OpenTAXIIPersistenceAPI
@@ -801,61 +802,44 @@ def _apply_match_version(
         if "all" in match_version:
             return query
         version_filters = []
+
         for value in match_version:
             if value == "first":
+                VersionSTIXObject = aliased(taxii2models.STIXObject, name="oso_min")
+
                 min_versions_subq = (
-                    self.db.session.query(
-                        taxii2models.STIXObject.id,
-                        func.min(taxii2models.STIXObject.version).label("min_version"),
+                    sqlalchemy.select(
+                        func.min(VersionSTIXObject.version).label("min_version"),
                     )
-                    .filter(
-                        taxii2models.STIXObject.collection_id == collection_id,
+                    .where(  # type: ignore[call-arg]
+                        VersionSTIXObject.collection_id == collection_id,
+                        VersionSTIXObject.id == taxii2models.STIXObject.id,
                     )
-                    .group_by(taxii2models.STIXObject.id)
-                    .subquery()
+                    .group_by(VersionSTIXObject.id)
+                    .scalar_subquery()  # type: ignore[attr-defined]
+                    .correlate(taxii2models.STIXObject)
                 )
-                min_version_pks = (
-                    self.db.session.query(taxii2models.STIXObject.pk)
-                    .select_from(taxii2models.STIXObject)
-                    .join(
-                        min_versions_subq,
-                        (
-                            (taxii2models.STIXObject.id == min_versions_subq.c.id)
-                            & (
-                                taxii2models.STIXObject.version
-                                == min_versions_subq.c.min_version
-                            )
-                        ),
-                    )
+                version_filters.append(
+                    min_versions_subq == taxii2models.STIXObject.version
                 )
-                version_filters.append(taxii2models.STIXObject.pk.in_(min_version_pks))
             elif value == "last":
+                VersionSTIXObject = aliased(taxii2models.STIXObject, name="oso_max")
+
                 max_versions_subq = (
-                    self.db.session.query(
-                        taxii2models.STIXObject.id,
-                        func.max(taxii2models.STIXObject.version).label("max_version"),
+                    sqlalchemy.select(
+                        func.max(VersionSTIXObject.version).label("max_version")
                     )
-                    .filter(
-                        taxii2models.STIXObject.collection_id == collection_id,
+                    .where(  # type: ignore[call-arg]
+                        VersionSTIXObject.collection_id == collection_id,
+                        VersionSTIXObject.id == taxii2models.STIXObject.id,
                     )
-                    .group_by(taxii2models.STIXObject.id)
-                    .subquery()
+                    .group_by(VersionSTIXObject.id)
+                    .scalar_subquery()  # type: ignore[attr-defined]
+                    .correlate(taxii2models.STIXObject)
                 )
-                max_version_pks = (
-                    self.db.session.query(taxii2models.STIXObject.pk)
-                    .select_from(taxii2models.STIXObject)
-                    .join(
-                        max_versions_subq,
-                        (
-                            (taxii2models.STIXObject.id == max_versions_subq.c.id)
-                            & (
-                                taxii2models.STIXObject.version
-                                == max_versions_subq.c.max_version
-                            )
-                        ),
-                    )
+                version_filters.append(
+                    max_versions_subq == taxii2models.STIXObject.version
                 )
-                version_filters.append(taxii2models.STIXObject.pk.in_(max_version_pks))
             else:
                 version_filters.append(taxii2models.STIXObject.version == value)
         query = query.filter(reduce(or_, version_filters))
@@ -874,7 +858,7 @@ def _apply_limit(
         self, query: Query, limit: Optional[int] = None
     ) -> Tuple[Query, bool]:
         if limit is not None:
-            more = limit < query.count()
+            more = limit < query.limit(limit + 1).count()
             query = query.limit(limit)
         else:
             more = False

opentaxii/persistence/sqldb/taxii2models.py

@@ -163,7 +163,9 @@ class STIXObject(Base):
 
     __table_args__ = (
         sqlalchemy.UniqueConstraint(collection_id, id, version),
-        sqlalchemy.Index("ix_opentaxii_stixobject_date_added_id", date_added, id),
+        sqlalchemy.Index(
+            "ix_opentaxii_stixobject_col_date_added_id", collection_id, date_added, id
+        ),
     )
 
     @classmethod