From 4bc31725b2e03a61397c474a77beb097418cf520 Mon Sep 17 00:00:00 2001
From: Maxime Petazzoni <maxime.petazzoni@bulix.org>
Date: Tue, 16 Apr 2024 11:39:43 -0700
Subject: [PATCH] feat: close the adapter's connection on quit (#518)

This change introduces a new optional method of HarlequinConnection
implementations in adapters, giving them an opportunity to cleanly close
the connection to the underlying database when the app exists.

The `close()` method is called in the `action_quit()` handler.
---
 CHANGELOG.md             | 4 ++++
 src/harlequin/adapter.py | 9 +++++++++
 src/harlequin/app.py     | 2 ++
 3 files changed, 15 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 457a00fa..a9b7d084 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file.
 
 ## [Unreleased]
 
+### Features
+
+-   A new `HarlequinConnection.close()` method can be implemented by adapters to gracefully close database connections when the application exits.
+
 ### Bug Fixes
 
 -   Wrong link on text_area_clipboard_error error message ([#509](https://github.com/tconbeer/harlequin/issues/509))
diff --git a/src/harlequin/adapter.py b/src/harlequin/adapter.py
index f926c916..9357463e 100644
--- a/src/harlequin/adapter.py
+++ b/src/harlequin/adapter.py
@@ -149,6 +149,15 @@ def validate_sql(self, text: str) -> str:
         """
         raise NotImplementedError
 
+    def close(self) -> None:
+        """
+        Closes the connection, if necessary. This function is called when the app
+        quits.
+
+        Returns: None
+        """
+        return None
+
 
 class HarlequinAdapter(ABC):
     """
diff --git a/src/harlequin/app.py b/src/harlequin/app.py
index 52e62bb8..60c6a90e 100644
--- a/src/harlequin/app.py
+++ b/src/harlequin/app.py
@@ -648,6 +648,8 @@ async def action_quit(self) -> None:
             s3_tree=self.data_catalog.s3_tree,
             history=self.history,
         )
+        if self.connection:
+            self.connection.close()
         await super().action_quit()
 
     def action_show_help_screen(self) -> None: