feat: BI-6252 propagate schema op_mode into US entries; skip excess validation in file conn import (#1024)

* feat: BI-6252 propagate schema op_mode into US entries; skip excessive validation during file conn imports

* move OperationsMode into dl_constants

* fix more imports

* fix more imports x2

* include op_mode in update_object as well

Author: KonstantAnxiety

lib/dl_api_connector/dl_api_connector/api_schema/extras.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import enum
 from typing import (
     Optional,
     Sequence,
@@ -10,27 +9,7 @@
 
 import attr
 
-
-class OperationsMode(enum.Enum):
-    pass
-
-
-class CreateMode(OperationsMode):
-    create = enum.auto()
-    test = enum.auto()
-
-
-class ImportMode(OperationsMode):
-    create_from_import = enum.auto()
-
-
-class EditMode(OperationsMode):
-    edit = enum.auto()
-    test = enum.auto()
-
-
-class ExportMode(OperationsMode):
-    export = enum.auto()
+from dl_constants.enums import OperationsMode
 
 
 class SchemaKWArgs(TypedDict):

lib/dl_api_connector/dl_api_connector/api_schema/top_level.py

@@ -28,13 +28,15 @@
 from marshmallow import fields as ma_fields
 
 from dl_api_connector.api_schema.extras import (
+    FieldExtra,
+    SchemaKWArgs,
+)
+from dl_constants.enums import (
     CreateMode,
     EditMode,
     ExportMode,
-    FieldExtra,
     ImportMode,
     OperationsMode,
-    SchemaKWArgs,
 )
 from dl_core import exc as bi_core_exc
 from dl_core.base_models import (
@@ -355,13 +357,16 @@ def create_object(self, data: dict[str, Any]) -> _US_ENTRY_TV:
 
         obj = self.TARGET_CLS.create_from_dict(  # type: ignore  # TODO: fix
             data_dict=self.create_data_model(data_attributes),
+            entry_op_mode=self.operations_mode,
             **self.default_create_from_dict_kwargs(data),
         )
         self.validate_object(obj)
         return obj
 
     @final
     def update_object(self, obj: _US_ENTRY_TV, data: dict[str, Any]) -> _US_ENTRY_TV:
+        obj.entry_op_mode = self.operations_mode
+
         # Assumed that only data of USEntry can be modified with schema
         assert not (data.keys() - {"data"})
 

lib/dl_api_connector/dl_api_connector_tests/unit/test_schemas_base.py

@@ -12,12 +12,12 @@
 )
 import pytest
 
-from dl_api_connector.api_schema.extras import (
+from dl_api_connector.api_schema.extras import FieldExtra
+from dl_api_connector.api_schema.top_level import BaseTopLevelSchema
+from dl_constants.enums import (
     CreateMode,
     EditMode,
-    FieldExtra,
 )
-from dl_api_connector.api_schema.top_level import BaseTopLevelSchema
 from dl_testing.utils import get_log_record
 
 

lib/dl_api_lib/dl_api_lib/app/control_api/resources/base.py

@@ -14,10 +14,10 @@
 from dl_api_commons.base_models import RequestContextInfo
 from dl_api_commons.flask.middlewares.commit_rci_middleware import ReqCtxInfoMiddleware
 from dl_api_commons.flask.required_resources import RequiredResourceCommon
-from dl_api_connector.api_schema.extras import OperationsMode
 from dl_api_lib import api_decorators
 from dl_api_lib.schemas.tools import prepare_schema_context
 from dl_api_lib.service_registry.service_registry import ApiServiceRegistry
+from dl_constants.enums import OperationsMode
 from dl_core.flask_utils.us_manager_middleware import USManagerFlaskMiddleware
 from dl_core.us_manager.us_manager_sync import SyncUSManager
 

lib/dl_api_lib/dl_api_lib/app/control_api/resources/connections.py

@@ -15,12 +15,6 @@
 from dl_api_commons.flask.middlewares.logging_context import put_to_request_context
 from dl_api_commons.flask.required_resources import RequiredResourceCommon
 from dl_api_connector.api_schema.connection_base import ConnectionOptionsSchema
-from dl_api_connector.api_schema.extras import (
-    CreateMode,
-    EditMode,
-    ExportMode,
-    ImportMode,
-)
 from dl_api_lib import exc
 from dl_api_lib.api_decorators import schematic_request
 from dl_api_lib.app.control_api.resources import API
@@ -38,7 +32,13 @@
 )
 from dl_api_lib.schemas.main import ImportResponseSchema
 from dl_api_lib.utils import need_permission_on_entry
-from dl_constants.enums import ConnectionType
+from dl_constants.enums import (
+    ConnectionType,
+    CreateMode,
+    EditMode,
+    ExportMode,
+    ImportMode,
+)
 from dl_constants.exc import DLBaseException
 from dl_core.data_source.type_mapping import get_data_source_class
 from dl_core.data_source_merge_tools import make_spec_from_dict

lib/dl_api_lib/dl_api_lib/schemas/tools.py

@@ -6,11 +6,11 @@
     Optional,
 )
 
-from dl_api_connector.api_schema.extras import OperationsMode
 from dl_api_connector.api_schema.top_level import (
     BaseTopLevelSchema,
     USEntryBaseSchema,
 )
+from dl_constants.enums import OperationsMode
 from dl_core.us_manager.us_manager import USManagerBase
 
 

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/api/api_schema/source.py

@@ -17,17 +17,15 @@
     pre_load,
 )
 
-from dl_api_connector.api_schema.extras import (
-    CreateMode,
-    FieldExtra,
-    ImportMode,
-)
+from dl_api_connector.api_schema.extras import FieldExtra
 from dl_api_connector.api_schema.source_base import (
     SQLDataSourceSchema,
     SQLDataSourceTemplateSchema,
 )
 from dl_constants.enums import (
+    CreateMode,
     FileProcessingStatus,
+    ImportMode,
     UserDataType,
 )
 from dl_model_tools.schema.base import BaseSchema

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/core/us_connection.py

@@ -16,6 +16,7 @@
 from dl_constants.enums import (
     DataSourceRole,
     FileProcessingStatus,
+    ImportMode,
 )
 from dl_core import (
     connection_models,
@@ -261,6 +262,28 @@ def validate(self) -> None:
 
     _saved_sources: Optional[list[FileDataSource]] = None
 
+    async def _update_added_sources(self, sources_to_add: set[str]) -> None:
+        """enriches freshly added or replaced sources with info from respective DataFiles from file-uploader-api"""
+
+        fu_client_factory = self.us_manager.get_services_registry().get_file_uploader_client_factory()
+
+        sources_desc = []
+        for src_id in sources_to_add:
+            src = self.get_file_source_by_id(src_id)
+            sources_desc.append(src.get_desc())
+
+        if sources_desc:
+            async with fu_client_factory.get_client() as fu_client:
+                internal_params = await fu_client.get_internal_params_batch(sources_desc)
+
+            for src_order, src_id in enumerate(sources_to_add):
+                self.update_data_source(
+                    src_id,
+                    role=DataSourceRole.origin,
+                    raw_schema=internal_params[src_order].raw_schema,
+                    preview_id=internal_params[src_order].preview_id,
+                )
+
     async def validate_new_data(
         self,
         services_registry: ServicesRegistry,
@@ -269,7 +292,7 @@ async def validate_new_data(
     ) -> None:
         assert isinstance(original_version, (type(self), type(None)))
         if original_version is None:
-            saved_sources = set()
+            saved_sources: set[str] = set()
         else:
             saved_sources = set(src.id for src in original_version.data.sources)
 
@@ -303,28 +326,16 @@ async def validate_new_data(
                 # Restore original internal source properties
                 self.restore_source_params_from_orig(src_id, original_version)
 
-        fu_client_factory = self.us_manager.get_services_registry().get_file_uploader_client_factory()
-
-        sources_desc = []
-        for src_id in sources_to_add:
-            src = self.get_file_source_by_id(src_id)
-            sources_desc.append(src.get_desc())
-
-        if sources_desc:
-            async with fu_client_factory.get_client() as fu_client:
-                internal_params = await fu_client.get_internal_params_batch(sources_desc)
+        if self.entry_op_mode == ImportMode.create_from_import:
+            # nothing else to do for an imported connection, because original `DataFile`s most certainly don't exist
+            return
 
-            for src_order, src_id in enumerate(sources_to_add):
-                self.update_data_source(
-                    src_id,
-                    role=DataSourceRole.origin,
-                    raw_schema=internal_params[src_order].raw_schema,
-                    preview_id=internal_params[src_order].preview_id,
-                )
+        await self._update_added_sources(sources_to_add)
 
+        # trigger re-save for sources, that are marked for update, by removing them from _saved_sources
         if self._saved_sources is not None:
             saved_source_ids = set(src.id for src in self._saved_sources)
-            for src_id in sources_to_update & saved_source_ids:  # trigger update by removing source from _saved_sources
+            for src_id in sources_to_update & saved_source_ids:
                 self._remove_saved_source_by_id(src_id)
 
         # clear errors for updated sources

lib/dl_constants/dl_constants/enums.py

@@ -449,3 +449,25 @@ class MigrationStatus(Enum):
     migrated_up = "migrated_up"
     migrated_down = "migrated_down"
     error = "error"
+
+
+class OperationsMode(Enum):
+    pass
+
+
+class CreateMode(OperationsMode):
+    create = auto()
+    test = auto()
+
+
+class ImportMode(OperationsMode):
+    create_from_import = auto()
+
+
+class EditMode(OperationsMode):
+    edit = auto()
+    test = auto()
+
+
+class ExportMode(OperationsMode):
+    export = auto()

lib/dl_core/dl_core/us_connection_base.py

@@ -32,6 +32,7 @@
     DataSourceType,
     MigrationStatus,
     NotificationLevel,
+    OperationsMode,
     RawSQLLevel,
     UserDataType,
     is_raw_sql_level_dashsql_allowed,
@@ -200,6 +201,7 @@ def __init__(
         hidden: bool = False,
         data_strict: bool = True,
         migration_status: MigrationStatus = MigrationStatus.non_migrated,
+        entry_op_mode: Optional[OperationsMode] = None,
         *,
         us_manager: USManagerBase,
     ):
@@ -218,6 +220,7 @@ def __init__(
             hidden=hidden,
             data_strict=data_strict,
             migration_status=migration_status,
+            entry_op_mode=entry_op_mode,
             us_manager=us_manager,
         )