From 8ac232da13dde6cfa6671d84e861f95551a43cd5 Mon Sep 17 00:00:00 2001
From: KonstantAnxiety <konstasa@yandex-team.ru>
Date: Thu, 20 Feb 2025 13:47:54 +0300
Subject: [PATCH] fix(datasets): BI-6085 load dataset connection before
 updating it from body (control-api) (#847)

---
 .../dl_api_lib/api_common/dataset_loader.py     |  2 +-
 .../app/control_api/resources/dataset.py        |  6 ++----
 .../app/control_api/resources/dataset_base.py   | 17 +++++++++++------
 3 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/lib/dl_api_lib/dl_api_lib/api_common/dataset_loader.py b/lib/dl_api_lib/dl_api_lib/api_common/dataset_loader.py
index 6f02ef40f..fccddbd6f 100644
--- a/lib/dl_api_lib/dl_api_lib/api_common/dataset_loader.py
+++ b/lib/dl_api_lib/dl_api_lib/api_common/dataset_loader.py
@@ -153,7 +153,7 @@ def _update_dataset_sources_from_body(
                     connection_id=connection_id,
                     source_type=source_type,
                     **source_data["parameters"],
-                )  # not that this does not include title and raw_schema updates
+                )  # note that this does not include title and raw_schema updates
                 old_raw_schema = old_src_coll.get_cached_raw_schema(role=DataSourceRole.origin)
                 schema_updated = not are_raw_schemas_same(old_raw_schema, source_data["raw_schema"])
 
diff --git a/lib/dl_api_lib/dl_api_lib/app/control_api/resources/dataset.py b/lib/dl_api_lib/dl_api_lib/app/control_api/resources/dataset.py
index 4775108d8..355e77c96 100644
--- a/lib/dl_api_lib/dl_api_lib/app/control_api/resources/dataset.py
+++ b/lib/dl_api_lib/dl_api_lib/app/control_api/resources/dataset.py
@@ -124,7 +124,7 @@ class DatasetItem(BIResource):
     def delete(self, dataset_id):
         """Delete dataset"""
         us_manager = self.get_us_manager()
-        ds, _ = DatasetResource.get_dataset(dataset_id=dataset_id, body={})
+        ds, _ = DatasetResource.get_dataset(dataset_id=dataset_id, body={}, load_dependencies=False)
         utils.need_permission_on_entry(ds, USPermissionKind.admin)
 
         us_manager.delete(ds)
@@ -141,7 +141,7 @@ class DatasetItemFields(BIResource):
         },
     )
     def get(self, dataset_id):  # type: ignore  # TODO: fix
-        ds, _ = DatasetResource.get_dataset(dataset_id=dataset_id, body={})
+        ds, _ = DatasetResource.get_dataset(dataset_id=dataset_id, body={}, load_dependencies=False)
         fields = [
             {
                 "title": f.title,
@@ -172,7 +172,6 @@ def post(self, dataset_id, body):  # type: ignore  # TODO: fix
         copy_us_key = body["new_key"]
         us_manager = self.get_us_manager()
         ds, _ = self.get_dataset(dataset_id=dataset_id, body={})
-        us_manager.load_dependencies(ds)
         orig_ds_loc = ds.entry_key
         copy_ds_loc: PathEntryLocation
 
@@ -216,7 +215,6 @@ def get(self, dataset_id, version):  # type: ignore  # TODO: fix
 
         ds_dict["is_favorite"] = ds.is_favorite
 
-        us_manager.load_dependencies(ds)
         ds_dict.update(self.make_dataset_response_data(dataset=ds, us_entry_buffer=us_manager.get_entry_buffer()))
         return ds_dict
 
diff --git a/lib/dl_api_lib/dl_api_lib/app/control_api/resources/dataset_base.py b/lib/dl_api_lib/dl_api_lib/app/control_api/resources/dataset_base.py
index d76b13282..caaf1b996 100644
--- a/lib/dl_api_lib/dl_api_lib/app/control_api/resources/dataset_base.py
+++ b/lib/dl_api_lib/dl_api_lib/app/control_api/resources/dataset_base.py
@@ -4,10 +4,7 @@
 import logging
 from typing import (
     Any,
-    List,
     Optional,
-    Set,
-    Tuple,
 )
 
 from dl_api_lib.api_common.dataset_loader import (
@@ -83,7 +80,12 @@ def create_dataset_api_loader(cls) -> DatasetApiLoader:
         )
 
     @classmethod
-    def get_dataset(cls, dataset_id: Optional[str], body: dict) -> Tuple[Dataset, DatasetUpdateInfo]:
+    def get_dataset(
+        cls,
+        dataset_id: Optional[str],
+        body: dict,
+        load_dependencies: bool = True,
+    ) -> tuple[Dataset, DatasetUpdateInfo]:
         us_manager = cls.get_us_manager()
         if dataset_id:
             try:
@@ -96,6 +98,9 @@ def get_dataset(cls, dataset_id: Optional[str], body: dict) -> Tuple[Dataset, Da
                 us_manager=us_manager,
             )
 
+        if load_dependencies:
+            us_manager.load_dependencies(dataset)
+
         loader = cls.create_dataset_api_loader()
         update_info = loader.update_dataset_from_body(
             dataset=dataset,
@@ -254,7 +259,7 @@ def dump_option_data(
 
         opt_data["sources"] = dict(items=[])
         connection_ids = set()
-        connection_types: Set[Optional[ConnectionType]] = set()
+        connection_types: set[Optional[ConnectionType]] = set()
         for source_id in ds_accessor.get_data_source_id_list():
             dsrc_coll_spec = ds_accessor.get_data_source_coll_spec_strict(source_id=source_id)
             dsrc_coll = dsrc_coll_factory.get_data_source_collection(spec=dsrc_coll_spec)
@@ -297,7 +302,7 @@ def dump_option_data(
                 )
             )
 
-        compatible_conn_types: List[dict] = []
+        compatible_conn_types: list[dict] = []
         for conn_type in capabilities.get_compatible_connection_types():
             if connection_ids:  # There already are connections in the dataset
                 continue