cancel python model job when dbt exit

dbt/adapters/databricks/auth.py

@@ -15,21 +15,23 @@
 
 class token_auth(CredentialsProvider):
     _token: str
+    _host: str
 
-    def __init__(self, token: str) -> None:
+    def __init__(self, token: str, host: str) -> None:
         self._token = token
+        self._host = host
 
     def auth_type(self) -> str:
         return "token"
 
     def as_dict(self) -> dict:
-        return {"token": self._token}
+        return {"token": self._token, "host": self._host}
 
     @staticmethod
     def from_dict(raw: Optional[dict]) -> Optional[CredentialsProvider]:
         if not raw:
             return None
-        return token_auth(raw["token"])
+        return token_auth(raw["token"], raw["host"])
 
     def __call__(self, _: Optional[Config] = None) -> HeaderFactory:
         static_credentials = {"Authorization": f"Bearer {self._token}"}

dbt/adapters/databricks/connections.py

@@ -77,6 +77,8 @@
 from dbt_common.utils import cast_to_str
 from requests import Session
 
+from dbt.adapters.databricks.python_submissions import BaseDatabricksHelper
+
 if TYPE_CHECKING:
     from agate import Table
 
@@ -475,6 +477,14 @@ class DatabricksConnectionManager(SparkConnectionManager):
     TYPE: str = "databricks"
     credentials_provider: Optional[TCredentialProvider] = None
 
+    def cancel_open(self) -> List[str]:
+        for run_id in BaseDatabricksHelper.run_ids:
+            logger.debug(f"Cancel python model job: {run_id}")
+            BaseDatabricksHelper.cancel_run_id(run_id, self.credentials_provider.as_dict()['token'], self.credentials_provider.as_dict()['host'])
+        BaseDatabricksHelper.run_ids.clear()
+        return super().cancel_open()
+
+
     def compare_dbr_version(self, major: int, minor: int) -> int:
         version = (major, minor)
 

dbt/adapters/databricks/credentials.py

@@ -255,7 +255,7 @@ def authenticate(self, in_provider: Optional[TCredentialProvider]) -> TCredentia
         self._lock.acquire()
         try:
             if self.token:
-                provider = token_auth(self.token)
+                provider = token_auth(self.token, self.host)
                 self._credentials_provider = provider.as_dict()
                 return provider
 

dbt/adapters/databricks/python_submissions.py

@@ -20,13 +20,19 @@
 from dbt.adapters.databricks.credentials import TCredentialProvider
 from dbt.adapters.databricks.logging import logger
 
+import threading
+
 
 DEFAULT_POLLING_INTERVAL = 10
 SUBMISSION_LANGUAGE = "python"
 DEFAULT_TIMEOUT = 60 * 60 * 24
 
 
 class BaseDatabricksHelper(PythonJobHelper):
+
+    run_ids = list()
+    _lock = threading.Lock()  # to avoid concurrent issue
+
     def __init__(self, parsed_model: Dict, credentials: DatabricksCredentials) -> None:
         self.credentials = credentials
         self.identifier = parsed_model["alias"]
@@ -45,6 +51,7 @@ def __init__(self, parsed_model: Dict, credentials: DatabricksCredentials) -> No
         self.extra_headers = {
             "User-Agent": f"dbt-databricks/{version}",
         }
+        host = credentials.host
 
     @property
     def cluster_id(self) -> str:
@@ -96,6 +103,7 @@ def _upload_notebook(self, path: str, compiled_code: str) -> None:
         if response.status_code != 200:
             raise DbtRuntimeError(f"Error creating python notebook.\n {response.content!r}")
 
+
     def _submit_job(self, path: str, cluster_spec: dict) -> str:
         job_spec = {
             "run_name": f"{self.schema}-{self.identifier}-{uuid.uuid4()}",
@@ -146,6 +154,7 @@ def _submit_through_notebook(self, compiled_code: str, cluster_spec: dict) -> No
 
         # submit job
         run_id = self._submit_job(whole_file_path, cluster_spec)
+        self._insert_run_ids(run_id)
 
         self.polling(
             status_func=self.session.get,
@@ -173,6 +182,51 @@ def _submit_through_notebook(self, compiled_code: str, cluster_spec: dict) -> No
                 "match the line number in your code due to dbt templating)\n"
                 f"{utils.remove_ansi(json_run_output['error_trace'])}"
             )
+        self._remove_run_ids(run_id)
+
+
+    def _remove_run_ids(self, run_id: str):
+        self._lock.acquire()
+        try:
+            self.run_ids.remove(run_id)
+        except Exception as e:
+            logger.warning(e)
+        finally:
+            self._lock.release()
+
+
+    def _insert_run_ids(self, run_id: str):
+        self._lock.acquire()
+        try:
+            self.run_ids.append(run_id)
+        except Exception as e:
+            logger.warning(e)
+        finally:
+            self._lock.release()
+
+    @staticmethod
+    def cancel_run_id(run_id: str, token: str, host: str) -> None:
+        retry_strategy = Retry(total=4, backoff_factor=0.5)
+        adapter = HTTPAdapter(max_retries=retry_strategy)
+        session = Session()
+        session.mount("https://", adapter)
+        extra_headers = {
+            "User-Agent": f"dbt-databricks/{version}",
+            "Authorization": f"Bearer {token}"
+        }
+
+        response = session.post(
+            f"https://{host}/api/2.0/jobs/runs/cancel",
+            headers=extra_headers,
+            json={
+                "run_id": run_id
+            },
+        )
+
+        print(response.status_code)
+        if response.status_code != 200:
+            logger.warning(f"Cancel run id failed.\n {response.content!r}")
+
 
     def submit(self, compiled_code: str) -> None:
         raise NotImplementedError(