|
60 | 60 | # Nullable types are problematic |
61 | 61 | "clickhouse", |
62 | 62 | "starrocks", |
| 63 | + # Db2 rejects table names starting with underscore (_versions, _snapshots, |
| 64 | + # _environments, _intervals). Use a separate state_connection (e.g., DuckDB). |
| 65 | + "db2", |
63 | 66 | } |
64 | 67 | MOTHERDUCK_TOKEN_REGEX = re.compile(r"(\?|\&)(motherduck_token=)(\S*)") |
65 | 68 | PASSWORD_REGEX = re.compile(r"(password=)(\S+)") |
@@ -2602,6 +2605,93 @@ def _connection_factory(self) -> t.Callable: |
2602 | 2605 | BaseDuckDBConnectionConfig, # type: ignore[type-abstract] |
2603 | 2606 | } |
2604 | 2607 |
|
| 2608 | + |
| 2609 | +class Db2ConnectionConfig(ConnectionConfig): |
| 2610 | + host: str |
| 2611 | + port: int = 50000 |
| 2612 | + database: str |
| 2613 | + db2_schema: str |
| 2614 | + username: str |
| 2615 | + password: str |
| 2616 | + ssl: bool = False |
| 2617 | + ssl_cert: t.Optional[str] = None |
| 2618 | + ssl_key: t.Optional[str] = None |
| 2619 | + ssl_ca: t.Optional[str] = None |
| 2620 | + connect_timeout: int = 30 |
| 2621 | + |
| 2622 | + concurrent_tasks: int = 4 |
| 2623 | + register_comments: bool = True |
| 2624 | + pre_ping: bool = True |
| 2625 | + |
| 2626 | + type_: t.Literal["db2"] = Field(alias="type", default="db2") |
| 2627 | + DIALECT: t.ClassVar[t.Literal["db2"]] = "db2" |
| 2628 | + DISPLAY_NAME: t.ClassVar[t.Literal["Db2"]] = "Db2" |
| 2629 | + DISPLAY_ORDER: t.ClassVar[t.Literal[19]] = 19 |
| 2630 | + |
| 2631 | + _engine_import_validator = _get_engine_import_validator("ibm_db", "db2") |
| 2632 | + |
| 2633 | + @property |
| 2634 | + def _connection_kwargs_keys(self) -> t.Set[str]: |
| 2635 | + return { |
| 2636 | + "host", |
| 2637 | + "port", |
| 2638 | + "database", |
| 2639 | + "db2_schema", |
| 2640 | + "username", |
| 2641 | + "password", |
| 2642 | + } |
| 2643 | + |
| 2644 | + @property |
| 2645 | + def _engine_adapter(self) -> t.Type[EngineAdapter]: |
| 2646 | + adapter = getattr(engine_adapter, "Db2EngineAdapter", None) |
| 2647 | + if adapter is None: |
| 2648 | + raise ConfigError( |
| 2649 | + "The Db2 engine adapter requires Python 3.10+ and the 'db2' extra. " |
| 2650 | + "Please install it with: pip install 'sqlmesh[db2]'" |
| 2651 | + ) |
| 2652 | + return t.cast(t.Type[EngineAdapter], adapter) |
| 2653 | + |
| 2654 | + def get_catalog(self) -> t.Optional[str]: |
| 2655 | + """Db2 stores catalog names in uppercase; normalise here so the default_catalog |
| 2656 | + passed to the adapter matches what get_current_catalog() returns at runtime.""" |
| 2657 | + catalog = super().get_catalog() |
| 2658 | + return catalog.upper() if catalog else None |
| 2659 | + |
| 2660 | + @property |
| 2661 | + def _connection_factory(self) -> t.Callable: |
| 2662 | + import ibm_db_dbi # type: ignore |
| 2663 | + |
| 2664 | + ssl = self.ssl |
| 2665 | + ssl_cert = self.ssl_cert |
| 2666 | + ssl_key = self.ssl_key |
| 2667 | + ssl_ca = self.ssl_ca |
| 2668 | + connect_timeout = self.connect_timeout |
| 2669 | + |
| 2670 | + def connect_db2(**kwargs: t.Any) -> t.Any: |
| 2671 | + conn_str_parts = [ |
| 2672 | + f"DATABASE={kwargs['database']}", |
| 2673 | + f"HOSTNAME={kwargs['host']}", |
| 2674 | + f"PORT={kwargs['port']}", |
| 2675 | + "PROTOCOL=TCPIP", |
| 2676 | + f"UID={kwargs['username']}", |
| 2677 | + f"PWD={kwargs['password']}", |
| 2678 | + f"CURRENTSCHEMA={kwargs['db2_schema']}", |
| 2679 | + f"CONNECTTIMEOUT={connect_timeout}", |
| 2680 | + ] |
| 2681 | + if ssl: |
| 2682 | + conn_str_parts.append("SECURITY=SSL") |
| 2683 | + if ssl_cert: |
| 2684 | + conn_str_parts.append(f"SSLClientCertificate={ssl_cert}") |
| 2685 | + if ssl_key: |
| 2686 | + conn_str_parts.append(f"SSLClientKey={ssl_key}") |
| 2687 | + if ssl_ca: |
| 2688 | + conn_str_parts.append(f"SSLServerCertificate={ssl_ca}") |
| 2689 | + conn_str = ";".join(conn_str_parts) + ";" |
| 2690 | + return ibm_db_dbi.connect(conn_str, "", "") |
| 2691 | + |
| 2692 | + return connect_db2 |
| 2693 | + |
| 2694 | + |
2605 | 2695 | CONNECTION_CONFIG_TO_TYPE = { |
2606 | 2696 | # Map all subclasses of ConnectionConfig to the value of their `type_` field. |
2607 | 2697 | tpe.all_field_infos()["type_"].default: tpe |
|
0 commit comments