BugFix

helviojunior · May 17, 2024 · 1fe73b6 · 1fe73b6
1 parent 97da42f
commit 1fe73b6
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/filecrawler/libs/database.py b/filecrawler/libs/database.py
@@ -64,6 +64,14 @@ def __init__(self, auto_create=True, db_name=None):
         else:
             self.connect_to_db()
 
+    def __enter__(self):
+        # make a database connection and return it
+        return self.db_connection
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # make sure the dbconnection gets closed
+        self.db_connection.close()
+
     @connect
     def insert_one(self, conn: Connection, table_name, **kwargs):
         table_name = self.scrub(table_name)