From 6b271709e7104e2be1084988a37a03d027d0c32d Mon Sep 17 00:00:00 2001
From: Tasko Olevski <tasko.olevski@sdsc.ethz.ch>
Date: Sat, 26 Oct 2024 21:40:12 -0400
Subject: [PATCH] squashme: use relationships to get data connectors

---
 .../renku_data_services/data_connectors/db.py | 54 ++++++-------------
 .../data_connectors/orm.py                    |  2 +
 2 files changed, 18 insertions(+), 38 deletions(-)

diff --git a/components/renku_data_services/data_connectors/db.py b/components/renku_data_services/data_connectors/db.py
index 138738a12..3d3541352 100644
--- a/components/renku_data_services/data_connectors/db.py
+++ b/components/renku_data_services/data_connectors/db.py
@@ -1,7 +1,7 @@
 """Adapters for data connectors database classes."""
 
 from collections.abc import AsyncIterator, Callable
-from typing import TypeVar, cast
+from typing import TypeVar
 
 from cryptography.hazmat.primitives.asymmetric import rsa
 from sqlalchemy import Select, delete, func, or_, select
@@ -503,49 +503,27 @@ async def get_data_connectors_with_secrets(
 
         async with self.session_maker() as session:
             stmt = (
-                select(schemas.DataConnectorORM, schemas.DataConnectorSecretORM)
-                .select_from(schemas.DataConnectorORM)  # NOTE: Makes sure the FROM statement is as expected
-                .join(
-                    target=schemas.DataConnectorToProjectLinkORM,
-                    onclause=schemas.DataConnectorORM.id == schemas.DataConnectorToProjectLinkORM.data_connector_id,
-                )
-                .join(
-                    target=schemas.DataConnectorSecretORM,
-                    onclause=schemas.DataConnectorORM.id == schemas.DataConnectorSecretORM.data_connector_id,
-                    isouter=True,  # NOTE: enables us to select data connectors with and without secrets
+                select(schemas.DataConnectorORM)
+                .where(
+                    schemas.DataConnectorORM.project_links.any(
+                        schemas.DataConnectorToProjectLinkORM.project_id == project_id
+                    )
                 )
-                .where(schemas.DataConnectorToProjectLinkORM.project_id == project_id)
                 .where(
                     or_(
-                        schemas.DataConnectorSecretORM.user_id == user.id,
-                        # NOTE: the user_id field on a connector secret is non-nullable, but
-                        # since we are doing an outer join this allows us to include data connectors
-                        # without secrets.
-                        schemas.DataConnectorSecretORM.user_id.is_(None),
+                        # Data connectors with secrets for the specific user
+                        schemas.DataConnectorORM.secrets.any(
+                            schemas.DataConnectorSecretORM.user_id == user.id,
+                        ),
+                        # Data connectors without any secrets
+                        # See: https://docs.sqlalchemy.org/en/20/orm/queryguide/select.html#exists-forms-has-any
+                        ~schemas.DataConnectorORM.secrets.any(),
                     )
                 )
-                # NOTE: The order is important for the processing of the data below
-                .order_by(schemas.DataConnectorORM.id)
-                .order_by(schemas.DataConnectorSecretORM.secret_id)
             )
-            results = await session.stream(stmt)
-            dc_current: models.DataConnector | None = None
-            dc_secrets: list[models.DataConnectorSecret] = []
-            async for res in results:
-                # NOTE: sqlalchemy does not set the types right for outer joins
-                dc, sec = cast(tuple[schemas.DataConnectorORM, schemas.DataConnectorSecretORM | None], res.t)
-                if dc_current is not None and dc.id != dc_current.id:
-                    yield models.DataConnectorWithSecrets(dc_current, dc_secrets)
-                if dc_current is None or dc.id != dc_current.id:
-                    dc_current = dc.dump()
-                    dc_secrets = [sec.dump()] if sec else []
-                    continue
-                if sec:
-                    dc_secrets.append(sec.dump())
-            if dc_current is None:
-                # There are no data connectors at all returned from the DB
-                return
-            yield models.DataConnectorWithSecrets(dc_current, dc_secrets)
+            results = await session.stream_scalars(stmt)
+            async for dc in results:
+                yield models.DataConnectorWithSecrets(dc.dump(), [secret.dump() for secret in dc.secrets])
 
     async def get_data_connector_secrets(
         self,
diff --git a/components/renku_data_services/data_connectors/orm.py b/components/renku_data_services/data_connectors/orm.py
index 447168d34..9992d647e 100644
--- a/components/renku_data_services/data_connectors/orm.py
+++ b/components/renku_data_services/data_connectors/orm.py
@@ -85,6 +85,8 @@ class DataConnectorORM(BaseORM):
         onupdate=func.now(),
         nullable=False,
     )
+    secrets: Mapped[list["DataConnectorSecretORM"]] = relationship(init=False, viewonly=True, lazy="selectin")
+    project_links: Mapped[list["DataConnectorToProjectLinkORM"]] = relationship(init=False, viewonly=True)
 
     def dump(self) -> models.DataConnector:
         """Create a data connector model from the DataConnectorORM."""