Revert previous dsrc migration fix (#137); implement a proper fix (#159)

datalens-tech · Dec 11, 2023 · fdb695c · fdb695c
1 parent 2ac053e
commit fdb695c
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 17 deletions.
diff --git a/lib/dl_connector_metrica/dl_connector_metrica/core/data_source_migration.py b/lib/dl_connector_metrica/dl_connector_metrica/core/data_source_migration.py
@@ -8,6 +8,7 @@
 
 class MetricaApiDataSourceMigrator(DefaultSQLDataSourceMigrator):
     table_source_type = SOURCE_TYPE_METRICA_API
+    with_db_name = True
 
 
 class AppMetricaApiDataSourceMigrator(DefaultSQLDataSourceMigrator):

diff --git a/lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql/data_source_migration.py b/lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql/data_source_migration.py
@@ -1,4 +1,5 @@
 from dl_core.connectors.sql_base.data_source_migration import DefaultSQLDataSourceMigrator
+from dl_core.data_source_spec.sql import StandardSchemaSQLDataSourceSpec
 
 from dl_connector_postgresql.core.postgresql.constants import (
     SOURCE_TYPE_PG_SUBSELECT,
@@ -8,5 +9,7 @@
 
 class PostgreSQLDataSourceMigrator(DefaultSQLDataSourceMigrator):
     table_source_type = SOURCE_TYPE_PG_TABLE
+    table_dsrc_spec_cls = StandardSchemaSQLDataSourceSpec
     subselect_source_type = SOURCE_TYPE_PG_SUBSELECT
+
     default_schema_name = "public"
diff --git a/lib/dl_core/dl_core/connectors/sql_base/data_source_migration.py b/lib/dl_core/dl_core/connectors/sql_base/data_source_migration.py
@@ -64,14 +64,15 @@ def _resolve_schema_name_for_import(
     ) -> Optional[str]:
         assert attr_name == "schema_name"
         schema_name: Optional[str] = getattr(migration_dto, attr_name, None)
-        if schema_name == self.default_schema_name:
-            schema_name = None
+        if schema_name is None:
+            schema_name = self.default_schema_name
         return schema_name
 
     def get_migration_specs(self) -> list[MigrationSpec]:
         result: list[MigrationSpec] = []
         if self.table_source_type is not None:
             assert self.table_dsrc_spec_cls is not None
+
             if self.with_db_name:
                 result.append(
                     MigrationSpec(
@@ -90,23 +91,23 @@ def get_migration_specs(self) -> list[MigrationSpec]:
                         ),
                     )
                 )
-            else:
-                result.append(
-                    MigrationSpec(
-                        source_type=self.table_source_type,
-                        dto_cls=SQLTableDSMI,
-                        dsrc_spec_cls=self.table_dsrc_spec_cls,
-                        migration_mapping_items=(
-                            MigrationKeyMappingItem(
-                                migration_dto_key="schema_name",
-                                source_spec_key="schema_name",
-                                custom_export_resolver=self._resolve_schema_name_for_export,
-                                custom_import_resolver=self._resolve_schema_name_for_import,
-                            ),
-                            MigrationKeyMappingItem(migration_dto_key="table_name", source_spec_key="table_name"),
+
+            result.append(
+                MigrationSpec(
+                    source_type=self.table_source_type,
+                    dto_cls=SQLTableDSMI,
+                    dsrc_spec_cls=self.table_dsrc_spec_cls,
+                    migration_mapping_items=(
+                        MigrationKeyMappingItem(
+                            migration_dto_key="schema_name",
+                            source_spec_key="schema_name",
+                            custom_export_resolver=self._resolve_schema_name_for_export,
+                            custom_import_resolver=self._resolve_schema_name_for_import,
                         ),
-                    )
+                        MigrationKeyMappingItem(migration_dto_key="table_name", source_spec_key="table_name"),
+                    ),
                 )
+            )
 
         if self.subselect_source_type is not None:
             assert self.subselect_dsrc_spec_cls is not None