11from abc import ABC
2+ from typing import Any
23from typing import Optional
34from typing import Tuple
45
56from databricks.sql.client import Connection
67from dbt.adapters.databricks.events.base import SQLErrorEvent
7- from dbt.contracts.graph.nodes import ResultNode
88
99
1010class ConnectionEvent(ABC):
11- def __init__(self, connection: Connection, message: str):
11+ def __init__(self, connection: Optional[ Connection] , message: str):
1212 self.message = message
1313 self.session_id = "Unknown"
1414 if connection:
@@ -19,31 +19,31 @@ def __str__(self) -> str:
1919
2020
2121class ConnectionCancel(ConnectionEvent):
22- def __init__(self, connection: Connection):
22+ def __init__(self, connection: Optional[ Connection] ):
2323 super().__init__(connection, "Cancelling connection")
2424
2525
2626class ConnectionClose(ConnectionEvent):
27- def __init__(self, connection: Connection):
27+ def __init__(self, connection: Optional[ Connection] ):
2828 super().__init__(connection, "Closing connection")
2929
3030
3131class ConnectionCancelError(ConnectionEvent):
32- def __init__(self, connection: Connection, exception: Exception):
32+ def __init__(self, connection: Optional[ Connection] , exception: Exception):
3333 super().__init__(
3434 connection, str(SQLErrorEvent(exception, "Exception while trying to cancel connection"))
3535 )
3636
3737
3838class ConnectionCloseError(ConnectionEvent):
39- def __init__(self, connection: Connection, exception: Exception):
39+ def __init__(self, connection: Optional[ Connection] , exception: Exception):
4040 super().__init__(
4141 connection, str(SQLErrorEvent(exception, "Exception while trying to close connection"))
4242 )
4343
4444
4545class ConnectionCreateError(ConnectionEvent):
46- def __init__(self, connection: Connection, exception: Exception):
46+ def __init__(self, connection: Optional[ Connection] , exception: Exception):
4747 super().__init__(
4848 connection, str(SQLErrorEvent(exception, "Exception while trying to create connection"))
4949 )
@@ -62,7 +62,7 @@ class ConnectionAcquire(ConnectionWrapperEvent):
6262 def __init__(
6363 self,
6464 description: str,
65- model: Optional[ResultNode ],
65+ model: Optional[Any ],
6666 compute_name: Optional[str],
6767 thread_identifier: Tuple[int, int],
6868 ):
0 commit comments