Added query compiler to all connector classes

datalens-tech · Oct 30, 2023 · c39d6ee · c39d6ee
1 parent a560f14
commit c39d6ee
Show file tree

Hide file tree

Showing 9 changed files with 38 additions and 1 deletion.
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
@@ -89,6 +89,8 @@ 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_bundle_chs3/dl_connector_bundle_chs3/chs3_base/core/connector.py b/lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/core/connector.py
@@ -13,6 +13,7 @@
 from dl_connector_bundle_chs3.chs3_base.core.constants import BACKEND_TYPE_CHS3
 from dl_connector_bundle_chs3.chs3_base.core.dto import BaseFileS3ConnDTO
 from dl_connector_bundle_chs3.chs3_base.core.type_transformer import FileTypeTransformer
+from dl_connector_clickhouse.core.clickhouse_base.query_compiler import ClickHouseQueryCompiler
 
 
 class BaseFileS3CoreConnectionDefinition(CoreConnectionDefinition):
@@ -30,3 +31,4 @@ class BaseFileS3CoreConnector(CoreConnector):
         {ConnSecuritySettings(NonUserInputConnectionSafetyChecker, frozenset({BaseFileS3ConnDTO}))}
     )
     query_cls = CHQuery
+    compiler_cls = ClickHouseQueryCompiler
diff --git a/lib/dl_connector_chyt/dl_connector_chyt/core/connector.py b/lib/dl_connector_chyt/dl_connector_chyt/core/connector.py
@@ -43,6 +43,7 @@
 )
 from dl_connector_chyt.core.type_transformer import CHYTTypeTransformer
 from dl_connector_chyt.core.us_connection import ConnectionCHYTToken
+from dl_connector_clickhouse.core.clickhouse_base.query_compiler import ClickHouseQueryCompiler
 
 
 class CHYTCoreConnectionDefinition(CoreConnectionDefinition):
@@ -101,3 +102,4 @@ class CHYTCoreConnector(CoreConnector):
         }
     )
     query_cls = CHQuery
+    compiler_cls = ClickHouseQueryCompiler
diff --git a/lib/dl_connector_greenplum/dl_connector_greenplum/core/connector.py b/lib/dl_connector_greenplum/dl_connector_greenplum/core/connector.py
@@ -26,6 +26,7 @@
     AsyncPostgresConnExecutor,
     PostgresConnExecutor,
 )
+from dl_connector_postgresql.core.postgresql_base.query_compiler import PostgreSQLQueryCompiler
 from dl_connector_postgresql.core.postgresql_base.sa_types import SQLALCHEMY_POSTGRES_TYPES
 from dl_connector_postgresql.core.postgresql_base.type_transformer import PostgreSQLTypeTransformer
 
@@ -60,3 +61,4 @@ class GreenplumCoreConnector(CoreConnector):
     )
     rqe_adapter_classes = frozenset({PostgresAdapter, AsyncPostgresAdapter})
     sa_types = SQLALCHEMY_POSTGRES_TYPES
+    compiler_cls = PostgreSQLQueryCompiler
diff --git a/lib/dl_connector_oracle/dl_connector_oracle/core/connector.py b/lib/dl_connector_oracle/dl_connector_oracle/core/connector.py
@@ -20,6 +20,7 @@
     OracleSubselectDataSource,
 )
 from dl_connector_oracle.core.data_source_migration import OracleDataSourceMigrator
+from dl_connector_oracle.core.query_compiler import OracleQueryCompiler
 from dl_connector_oracle.core.sa_types import SQLALCHEMY_ORACLE_TYPES
 from dl_connector_oracle.core.storage_schemas.connection import ConnectionSQLOracleDataStorageSchema
 from dl_connector_oracle.core.type_transformer import OracleServerTypeTransformer
@@ -56,3 +57,4 @@ class OracleCoreConnector(CoreConnector):
     )
     rqe_adapter_classes = frozenset({OracleDefaultAdapter})
     sa_types = SQLALCHEMY_ORACLE_TYPES
+    compiler_cls = OracleQueryCompiler
diff --git a/lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql/connector.py b/lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql/connector.py
@@ -26,6 +26,7 @@
     AsyncPostgresConnExecutor,
     PostgresConnExecutor,
 )
+from dl_connector_postgresql.core.postgresql_base.query_compiler import PostgreSQLQueryCompiler
 from dl_connector_postgresql.core.postgresql_base.sa_types import SQLALCHEMY_POSTGRES_TYPES
 from dl_connector_postgresql.core.postgresql_base.type_transformer import PostgreSQLTypeTransformer
 
@@ -60,3 +61,4 @@ class PostgreSQLCoreConnector(CoreConnector):
     )
     rqe_adapter_classes = frozenset({PostgresAdapter, AsyncPostgresAdapter})
     sa_types = SQLALCHEMY_POSTGRES_TYPES
+    compiler_cls = PostgreSQLQueryCompiler
diff --git a/lib/dl_core/dl_core/connectors/base/registrator.py b/lib/dl_core/dl_core/connectors/base/registrator.py
@@ -14,6 +14,7 @@
 from dl_core.connectors.base.dashsql import register_custom_dash_sql_key_names
 from dl_core.connectors.base.data_source_migration import register_data_source_migrator
 from dl_core.connectors.settings.registry import register_connector_settings_class
+from dl_core.data_processing.query_compiler_registry import register_sa_query_compiler_cls
 from dl_core.data_source.type_mapping import register_data_source_class
 from dl_core.data_source_spec.type_mapping import register_data_source_spec_class
 from dl_core.db.conversion_base import register_type_transformer_class
@@ -97,11 +98,15 @@ def register_connector(cls, connector: Type[CoreConnector]) -> None:
             conn_sec_settings.security_checker_cls.register_dto_types(dto_types)
         for adapter_cls in connector.rqe_adapter_classes:
             register_adapter_class(adapter_cls=adapter_cls)
-        register_sa_query_cls(backend_type=connector.backend_type, query_cls=connector.query_cls)
         register_query_fail_exceptions(exception_classes=connector.query_fail_exceptions)
         for notification_cls in connector.notification_classes:
             register_notification()(notification_cls)  # it is a parameterized decorator
 
+        # backend_type-dependent properties
+        backend_type = connector.backend_type
+        register_sa_query_cls(backend_type=backend_type, query_cls=connector.query_cls)
+        register_sa_query_compiler_cls(backend_type=backend_type, sa_query_compiler_cls=connector.compiler_cls)
+
         connector.registration_hook()  # for custom actions
 
 

diff --git a/lib/dl_core/dl_core/data_processing/query_compiler_registry.py b/lib/dl_core/dl_core/data_processing/query_compiler_registry.py
@@ -0,0 +1,18 @@
+from typing import Type
+
+from dl_constants.enums import SourceBackendType
+from dl_core.connectors.base.query_compiler import QueryCompiler
+
+
+_QUERY_COMPILER_CLS_BY_BACKEND: dict[SourceBackendType, Type[QueryCompiler]] = {}
+
+
+def register_sa_query_compiler_cls(backend_type: SourceBackendType, sa_query_compiler_cls: Type[QueryCompiler]) -> None:
+    if (registered_compiler_cls := _QUERY_COMPILER_CLS_BY_BACKEND.get(backend_type)) is not None:
+        assert registered_compiler_cls == sa_query_compiler_cls
+    else:
+        _QUERY_COMPILER_CLS_BY_BACKEND[backend_type] = sa_query_compiler_cls
+
+
+def get_sa_query_compiler_cls(backend_type: SourceBackendType) -> Type[QueryCompiler]:
+    return _QUERY_COMPILER_CLS_BY_BACKEND[backend_type]
diff --git a/lib/dl_core/dl_core/data_source/base.py b/lib/dl_core/dl_core/data_source/base.py
@@ -24,6 +24,7 @@
     ConnectionRef,
     SourceFilterSpec,
 )
+from dl_core.connectors.base.query_compiler import QueryCompiler
 from dl_core.data_source_spec.base import DataSourceSpec
 from dl_core.db import (
     IndexInfo,
@@ -70,6 +71,7 @@ class DataSource(metaclass=abc.ABCMeta):
     default_chunk_row_count: ClassVar[int] = 10000
     chunk_size_bytes: ClassVar[int] = 3 * 1024**2
 
+    compiler_cls: Type[QueryCompiler] = QueryCompiler
     conn_type: ClassVar[ConnectionType]  # TODO unbind DataSource and Connection classes BI-4083
 
     # TODO FIX: Remove ASAP