Fix clickhouse export (#3009)

* Hide all environment files

* remove prints

* fix asyncio issues

* improve doc for the hack

Author: ravenac95

.gitignore

@@ -32,6 +32,7 @@ data/
 .DS_Store
 *.pem
 *.env
+*.env.*
 .env*.local
 .env*.cloud
 *.log

warehouse/metrics_tools/transfer/clickhouse.py

@@ -48,7 +48,12 @@ async def import_table(
             loading_table_fqn,
             export_reference.columns.columns_as("clickhouse"),
         )
-        import_path = f"https://storage.googleapis.com/{gcs_bucket}/{gcs_blob_path}/*"
+        # We need the `2*` to match the files that are created by the export
+        # this is a bit of a hack due to some weird behavior in clickhouse that
+        # changed sometime around 2025-02-01. The `2*` is used because trino
+        # prefixes files with the date. So this will work for the next 975 years
+        # or so. Hopefully that's enough time.
+        import_path = f"https://storage.googleapis.com/{gcs_bucket}/{gcs_blob_path}/2*"
         self.logger.debug(f"Importing table {loading_table_fqn} from {gcs_path}")
         import_data(
             self.ch,

warehouse/oso_dagster/assets/clickhouse_dbt_marts.py

@@ -1,4 +1,5 @@
 import json
+import logging
 import os
 import uuid
 from pathlib import Path
@@ -16,6 +17,8 @@
 from ..utils.bq import BigQueryTableConfig
 from ..utils.common import SourceMode
 
+logger = logging.getLogger(__name__)
+
 MART_DIRECTORY = "marts"
 INTERMEDIATE_DIRECTORY = "intermediate"
 SYNC_KEY = "sync_to_db"
@@ -74,7 +77,7 @@ def all_clickhouse_dbt_mart_assets(
             table_name = n.get("name")
             # Only copy marts that are marked for sync
             if n.get("meta").get(SYNC_KEY, False):
-                print(f"Queuing {table_name}")
+                logger.debug(f"Queuing {table_name}")
                 copied_mart_names.append(table_name)
                 # Create an asset for each mart to copy
                 result = result + create_bq2clickhouse_asset(
@@ -131,8 +134,7 @@ def all_clickhouse_dbt_mart_assets(
             # Track which marts were skipped
             else:
                 skipped_mart_names.append(table_name)
-        print(
+        logger.debug(
             f"...queued {str(len(copied_mart_names))} marts, skipping {str(len(skipped_mart_names))}"
         )
-        # print(skipped_mart_names)
         return result

warehouse/oso_dagster/assets/dbt.py

@@ -1,3 +1,4 @@
+import logging
 import os
 from pathlib import Path
 from typing import Any, Dict, List, Mapping, Optional, Sequence
@@ -8,6 +9,8 @@
 
 from ..factories import AssetFactoryResponse, early_resources_asset_factory
 
+logger = logging.getLogger(__name__)
+
 
 class CustomDagsterDbtTranslator(DagsterDbtTranslator):
     def __init__(
@@ -81,8 +84,8 @@ def generate_dbt_asset(
     AssetsDefinition
         a single Dagster dbt asset
     """
-    print(f"Target[{target}] using profiles_dir({dbt_profiles_dir})")
-    print(f"\tmanifest_path({manifest_path})")
+    logger.debug(f"Target[{target}] using profiles_dir({dbt_profiles_dir})")
+    logger.debug(f"\tmanifest_path({manifest_path})")
     translator = CustomDagsterDbtTranslator(target, ["dbt", target], internal_map)
 
     asset_name = f"{target}_dbt"
@@ -94,7 +97,7 @@ def generate_dbt_asset(
         op_tags=op_tags,
     )
     def _generated_dbt_assets(context: AssetExecutionContext, config: DBTConfig):
-        print(f"using profiles dir {dbt_profiles_dir}")
+        logger.debug(f"using profiles dir {dbt_profiles_dir}")
         dbt = DbtCliResource(
             project_dir=os.fspath(project_dir),
             target=target,

warehouse/oso_dagster/cbt/bq.py

@@ -1,14 +1,18 @@
 # Query tools for bigquery tables
-from typing import List, cast, Optional, Dict, NotRequired, TypedDict
+import logging
 from functools import cache
+from typing import Dict, List, NotRequired, Optional, TypedDict, cast
 
 from google.cloud.bigquery import Client, Table, TableReference
 from google.cloud.bigquery.table import RowIterator
 from sqlglot import expressions as exp
-from .context import Connector, ColumnList
+
+from .context import ColumnList, Connector
 
 type ExtraVarType = str | int
 
+logger = logging.getLogger(__name__)
+
 
 class MetricQueryInput(TypedDict):
     ref: str
@@ -22,7 +26,7 @@ def __init__(self, bq: Client):
         self.bq = bq
 
     def __call__(self, table_ref: TableReference | Table | str):
-        print(table_ref)
+        logger.debug(table_ref)
         return BigQueryTableQueryHelper.load_by_table(self.bq, table_ref)
 
 
@@ -136,8 +140,6 @@ def update_columns_with(
 
         # If the other table is included then we can ensure only matching fields are contained
         if source_table is not None:
-            print("WHAT IS THIS")
-            print(source_table._table_ref)
             source_column_names = set(
                 map(lambda c: c.column_name, source_table.columns)
             )
@@ -147,10 +149,6 @@ def update_columns_with(
                 additional_columns = (
                     set(source_column_names) - set(ordered_columns) - set(exclude or [])
                 )
-                print("source")
-                print(list(source_column_names))
-                print("ordered")
-                print(list(ordered_columns))
                 if len(additional_columns) > 0:
                     raise Exception(
                         f"more columns in the source table than the destination `{additional_columns}`"

warehouse/oso_dagster/cbt/cbt.py

@@ -2,23 +2,20 @@
 #
 # A poor excuse for a dbt replacement when calling sql as a library
 import os
-import arrow
-from typing import List, Optional, Sequence
 from dataclasses import dataclass
 from enum import Enum
 from functools import cache
+from typing import List, Optional, Sequence
 
+import arrow
 from dagster import ConfigurableResource, DagsterLogManager
 from dagster_gcp import BigQueryResource
-from google.cloud.bigquery import (
-    TableReference,
-    Client,
-)
+from google.cloud.bigquery import Client, TableReference
 from google.cloud.exceptions import NotFound
 from jinja2 import Environment, FileSystemLoader, meta
 
-from .bq import BigQueryTableQueryHelper, BigQueryConnector
-from .context import DataContext, ContextQuery, Transformation
+from .bq import BigQueryConnector, BigQueryTableQueryHelper
+from .context import ContextQuery, DataContext, Transformation
 
 
 class UpdateStrategy(Enum):
@@ -311,7 +308,6 @@ def _transform_replace(
             select_query = context.transform_query(select_query, transformations).sql(
                 dialect="bigquery"
             )
-            print(select_query)
 
         if time_partitioning:
             self.log.debug("creating table with a time partition")

warehouse/oso_dagster/cbt/context.py

@@ -1,9 +1,12 @@
-from typing import Callable, TypeVar, List, Sequence, Tuple, Optional, cast
+import logging
 from collections import OrderedDict
+from typing import Callable, List, Optional, Sequence, Tuple, TypeVar, cast
 
 import sqlglot as sql
 from sqlglot import expressions as exp
 
+logger = logging.getLogger(__name__)
+
 
 T = TypeVar("T")
 ColumnList = List[Tuple[str, str]]
@@ -79,7 +82,7 @@ def execute_query(
         dialect: str = "bigquery",
     ) -> T:
         exp = self.transform_query(query, transformations, dialect=dialect)
-        print(exp.sql())
+        logger.debug(f"Executing query: {exp.sql()}")
         return self._connector.execute_expression(exp)
 
     # def examine_query(self, query: ContextQuery):

warehouse/oso_dagster/config.py

@@ -76,6 +76,8 @@ class DagsterConfig(BaseSettings):
     # `enable_k8s_executor` is used to enable k8s executor for dagster
     enable_k8s: bool = False
 
+    k8s_use_port_forward: bool = False
+
     trino_remote_url: str = "http://localhost:8080"
     trino_k8s_namespace: str = ""
     trino_k8s_service_name: str = ""

warehouse/oso_dagster/definitions.py

@@ -1,13 +1,14 @@
 import logging
 import os
+import warnings
 
-from dagster import Definitions
+from dagster import Definitions, ExperimentalWarning
 from dagster_dbt import DbtCliResource
 from dagster_embedded_elt.dlt import DagsterDltResource
 from dagster_gcp import BigQueryResource, GCSResource
 from dagster_k8s import k8s_job_executor
 from dagster_sqlmesh import SQLMeshContextConfig, SQLMeshResource
-from dotenv import load_dotenv
+from dotenv import find_dotenv, load_dotenv
 from metrics_tools.utils.logging import setup_module_logging
 from oso_dagster.resources.bq import BigQueryImporterResource
 from oso_dagster.resources.clickhouse import ClickhouseImporterResource
@@ -31,10 +32,10 @@
 from .resources import (
     BigQueryDataTransferResource,
     ClickhouseResource,
+    K8sApiResource,
     K8sResource,
     MCSK8sResource,
     MCSRemoteResource,
-    PodLocalK8sResource,
     PrefixedSQLMeshTranslator,
     Trino2BigQuerySQLMeshExporter,
     Trino2ClickhouseSQLMeshExporter,
@@ -54,10 +55,15 @@
 
 logger = logging.getLogger(__name__)
 
+if os.environ.get("ENV") == "local":
+    load_dotenv(find_dotenv(".env.local"))
+elif os.environ.get("ENV") == "production":
+    load_dotenv(find_dotenv(".env.production"))
 load_dotenv()
 
 
 def load_definitions():
+    warnings.filterwarnings("ignore", category=ExperimentalWarning)
     setup_module_logging("oso_dagster")
     # Load the configuration for the project
     global_config = DagsterConfig()  # type: ignore
@@ -127,19 +133,21 @@ def load_definitions():
 
     else:
         logger.info("Loading k8s resources")
-        k8s = PodLocalK8sResource()
+        k8s = K8sApiResource()
         trino = TrinoK8sResource(
             k8s=k8s,
             namespace=global_config.trino_k8s_namespace,
             service_name=global_config.trino_k8s_service_name,
             coordinator_deployment_name=global_config.trino_k8s_coordinator_deployment_name,
             worker_deployment_name=global_config.trino_k8s_worker_deployment_name,
+            use_port_forward=global_config.k8s_use_port_forward,
         )
         mcs = MCSK8sResource(
             k8s=k8s,
             namespace=global_config.mcs_k8s_namespace,
             service_name=global_config.mcs_k8s_service_name,
             deployment_name=global_config.mcs_k8s_deployment_name,
+            use_port_forward=global_config.k8s_use_port_forward,
         )
         sqlmesh_exporter = [
             Trino2ClickhouseSQLMeshExporter(

warehouse/oso_dagster/factories/dlt.py

@@ -42,7 +42,6 @@ def pydantic_to_dlt_nullable_columns(b: t.Type[BaseModel]):
     table_schema_columns = pydantic_to_table_schema_columns(b)
     for column in table_schema_columns.values():
         column["nullable"] = True
-    print(table_schema_columns)
     return table_schema_columns
 
 
@@ -123,9 +122,7 @@ def _factory(
                 # We need to ensure that both dlt and global_config are
                 # available to the generated asset as they're used by the
                 # generated function.
-                final_extra_resources = extra_resources.union(
-                    {"dlt", "global_config"}
-                )
+                final_extra_resources = extra_resources.union({"dlt", "global_config"})
 
                 @asset(
                     name=asset_name,
@@ -159,7 +156,9 @@ def _dlt_asset(
                     global_config = t.cast(
                         DagsterConfig, getattr(context.resources, "global_config")
                     )
-                    assert global_config, "global_config resource is not loading correctly"
+                    assert (
+                        global_config
+                    ), "global_config resource is not loading correctly"
                     if global_config.enable_bigquery:
                         context.log.debug("dlt pipeline setup to use staging")
                         pipeline = dltlib.pipeline(