dbt_adapter_skeleton_with_credentials_connection_manager_subclasses.py

1from dataclasses import dataclass
2from typing import Optional, List, Any
3
4from dbt.adapters.base import SQLAdapter
5from dbt.adapters.base.relation import BaseRelation
6from dbt.adapters.sql import SQLConnectionManager
7from dbt.contracts.connection import Connection, AdapterResponse
8from dbt_common.contracts.config.base import BaseConfig
9
10
11@dataclass
12class MyCredentials(BaseConfig):
13    """Define the parameters needed to connect to your database."""
14    host: str
15    port: int
16    user: str
17    password: str
18    database: str
19    schema: str
20
21    @property
22    def type(self):
23        return "my_adapter"
24
25    @property
26    def unique_field(self):
27        return self.host
28
29    def _connection_keys(self):
30        return ("host", "port", "user", "database", "schema")
31
32
33class MyConnectionManager(SQLConnectionManager):
34    """Handle the actual connection logic to the underlying database."""
35    TYPE = "my_adapter"
36
37    @classmethod
38    def open(cls, connection: Connection) -> Connection:
39        if connection.state == "open":
40            return connection
41        
42        credentials = connection.credentials
43        try:
44            # Replace with actual database driver connection logic
45            # handle = your_db_driver.connect(host=credentials.host, ...)
46            handle = None 
47            connection.handle = handle
48            connection.state = "open"
49        except Exception as e:
50            connection.state = "fail"
51            raise e
52        return connection
53
54    def cancel(self, connection: Connection):
55        connection.handle.close()
56
57    @classmethod
58    def get_response(cls, cursor: Any) -> AdapterResponse:
59        # Return a standardized response from the cursor
60        return AdapterResponse(_message="OK")
61
62
63class MyAdapter(SQLAdapter):
64    """The main adapter class bridging dbt and the database."""
65    ConnectionManager = MyConnectionManager
66    Relation = BaseRelation
67
68    @classmethod
69    def date_function(cls):
70        return "CURRENT_TIMESTAMP"
71
72    @classmethod
73    def is_cancelable(cls) -> bool:
74        return True