Removed query compiler from data sources

datalens-tech · Nov 8, 2023 · b4d1306 · b4d1306
1 parent afcf4bb
commit b4d1306
Show file tree

Hide file tree

Showing 10 changed files with 15 additions and 29 deletions.
diff --git a/lib/dl_connector_bigquery/dl_connector_bigquery/core/data_source.py b/lib/dl_connector_bigquery/dl_connector_bigquery/core/data_source.py
@@ -28,7 +28,6 @@
     BigQuerySubselectDataSourceSpec,
     BigQueryTableDataSourceSpec,
 )
-from dl_connector_bigquery.core.query_compiler import BigQueryQueryCompiler
 from dl_connector_bigquery.core.us_connection import ConnectionSQLBigQuery
 
 
@@ -45,8 +44,6 @@ class BigQueryTableDataSource(BigQueryDataSourceMixin, TableSQLDataSourceMixin,
     BigQuery table
     """
 
-    compiler_cls = BigQueryQueryCompiler
-
     @property
     def spec(self) -> BigQueryTableDataSourceSpec:
         assert isinstance(self._spec, BigQueryTableDataSourceSpec)
@@ -89,8 +86,6 @@ class BigQuerySubselectDataSource(BigQueryDataSourceMixin, SubselectDataSource):
     BigQuery subselect
     """
 
-    compiler_cls = BigQueryQueryCompiler
-
     @property
     def spec(self) -> BigQuerySubselectDataSourceSpec:
         assert isinstance(self._spec, BigQuerySubselectDataSourceSpec)

diff --git a/lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/data_source.py b/lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/data_source.py
@@ -15,7 +15,6 @@
 )
 
 from dl_connector_clickhouse.core.clickhouse_base.constants import CONNECTION_TYPE_CLICKHOUSE
-from dl_connector_clickhouse.core.clickhouse_base.query_compiler import ClickHouseQueryCompiler
 
 
 if TYPE_CHECKING:
@@ -45,7 +44,6 @@ class ClickHouseBaseMixin(BaseSQLDataSource):
             JoinType.right,
         }
     )
-    compiler_cls = ClickHouseQueryCompiler
 
     def get_connect_args(self) -> dict:
         return dict(super().get_connect_args(), server_version=self.db_version or self.default_server_version)

diff --git a/lib/dl_connector_greenplum/dl_connector_greenplum/core/data_source.py b/lib/dl_connector_greenplum/dl_connector_greenplum/core/data_source.py
@@ -12,12 +12,9 @@
     SOURCE_TYPE_GP_SUBSELECT,
     SOURCE_TYPE_GP_TABLE,
 )
-from dl_connector_postgresql.core.postgresql_base.query_compiler import PostgreSQLQueryCompiler
 
 
 class GreenplumDataSourceMixin(BaseSQLDataSource):
-    compiler_cls = PostgreSQLQueryCompiler
-
     conn_type = CONNECTION_TYPE_GREENPLUM
 
     @classmethod

diff --git a/lib/dl_connector_mysql/dl_connector_mysql/core/data_source.py b/lib/dl_connector_mysql/dl_connector_mysql/core/data_source.py
@@ -21,7 +21,6 @@
     SOURCE_TYPE_MYSQL_SUBSELECT,
     SOURCE_TYPE_MYSQL_TABLE,
 )
-from dl_connector_mysql.core.query_compiler import MySQLQueryCompiler
 
 
 LOGGER = logging.getLogger(__name__)
@@ -37,7 +36,6 @@ class MySQLDataSourceMixin(BaseSQLDataSource):
     )
 
     conn_type = CONNECTION_TYPE_MYSQL
-    compiler_cls = MySQLQueryCompiler
 
     @classmethod
     def is_compatible_with_type(cls, source_type: DataSourceType) -> bool:

diff --git a/lib/dl_connector_oracle/dl_connector_oracle/core/data_source.py b/lib/dl_connector_oracle/dl_connector_oracle/core/data_source.py
@@ -20,15 +20,12 @@
     SOURCE_TYPE_ORACLE_SUBSELECT,
     SOURCE_TYPE_ORACLE_TABLE,
 )
-from dl_connector_oracle.core.query_compiler import OracleQueryCompiler
 
 
 LOGGER = logging.getLogger(__name__)
 
 
 class OracleDataSourceMixin(BaseSQLDataSource):
-    compiler_cls = OracleQueryCompiler
-
     conn_type = CONNECTION_TYPE_ORACLE
 
     @classmethod

diff --git a/lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql/data_source.py b/lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql/data_source.py
@@ -20,11 +20,9 @@
     SOURCE_TYPE_PG_SUBSELECT,
     SOURCE_TYPE_PG_TABLE,
 )
-from dl_connector_postgresql.core.postgresql_base.query_compiler import PostgreSQLQueryCompiler
 
 
 class PostgreSQLDataSourceMixin(BaseSQLDataSource):
-    compiler_cls = PostgreSQLQueryCompiler
     supported_join_types: ClassVar[FrozenSet[JoinType]] = frozenset(
         {
             JoinType.inner,

diff --git a/lib/dl_core/dl_core/data_processing/prepared_components/default_manager.py b/lib/dl_core/dl_core/data_processing/prepared_components/default_manager.py
@@ -11,11 +11,13 @@
 from sqlalchemy.sql.selectable import FromClause
 
 from dl_constants.enums import DataSourceRole
+from dl_core.backend_types import get_backend_type
 from dl_core.components.accessor import DatasetComponentAccessor
 from dl_core.components.ids import AvatarId
 from dl_core.constants import DataAPILimits
 from dl_core.data_processing.prepared_components.manager_base import PreparedComponentManagerBase
 from dl_core.data_processing.prepared_components.primitives import PreparedSingleFromInfo
+from dl_core.data_processing.query_compiler_registry import get_sa_query_compiler_cls
 from dl_core.data_source.collection import DataSourceCollectionFactory
 import dl_core.data_source.sql
 import dl_core.exc as exc
@@ -80,7 +82,12 @@ def get_columns():  # type: ignore  # TODO: fix
         columns = get_columns()
         col_names = [col.name for col in columns]
         from_subquery = from_subquery and dsrc.supports_preview_from_subquery
-        query_compiler = dsrc.get_query_compiler()
+        connection = dsrc.connection
+        conn_type = connection.conn_type
+        backend_type = get_backend_type(conn_type=conn_type)
+        sa_dialect = connection.get_dialect()
+        query_compiler_cls = get_sa_query_compiler_cls(backend_type=backend_type)
+        query_compiler = query_compiler_cls(dialect=sa_dialect)
 
         sql_source: FromClause
         if from_subquery:

diff --git a/lib/dl_core/dl_core/data_processing/selectors/db.py b/lib/dl_core/dl_core/data_processing/selectors/db.py
@@ -29,6 +29,7 @@
 @attr.s
 class DatasetDbDataSelectorAsync(DatasetDataSelectorAsyncBase):
     """Async selector that fetches data from the database"""
+
     # TODO: Merge all selector logic into data processors
 
     _active_queries: list[BIQueryExecutionContext] = attr.ib(init=False, factory=list)

diff --git a/lib/dl_core/dl_core/data_source/sql.py b/lib/dl_core/dl_core/data_source/sql.py
@@ -147,11 +147,6 @@ def get_schema_info(self, conn_executor_factory: Callable[[], SyncConnExecutorBa
         schema_info = conn_executor.get_table_schema_info(table_def)
         return self._postprocess_raw_schema_from_db(schema_info)
 
-    def get_query_compiler(self) -> QueryCompiler:
-        return self.compiler_cls(
-            dialect=self.get_dialect(),
-        )
-
     def _get_db_version(self, conn_executor_factory: Callable[[], SyncConnExecutorBase]) -> Optional[str]:
         conn_executor = conn_executor_factory()
         return conn_executor.get_db_version(

diff --git a/lib/dl_core/dl_core/services_registry/selector_factory.py b/lib/dl_core/dl_core/services_registry/selector_factory.py
@@ -103,9 +103,9 @@ def _create_dataset_selector(
         allow_cache_usage: bool = True,
     ) -> DatasetDataSelectorAsyncBase:
         return DatasetDbDataSelectorAsync(
-                dataset=dataset,
-                service_registry=self.services_registry,
-                # allow_cache_usage=allow_cache_usage,
-                # is_bleeding_edge_user=self._is_bleeding_edge_user,
-                us_entry_buffer=us_entry_buffer,
-            )
+            dataset=dataset,
+            service_registry=self.services_registry,
+            # allow_cache_usage=allow_cache_usage,
+            # is_bleeding_edge_user=self._is_bleeding_edge_user,
+            us_entry_buffer=us_entry_buffer,
+        )