Skip to content

Commit

Permalink
Have the local environment release the connection when it's not in us…
Browse files Browse the repository at this point in the history
…e and not necessary for in-memory runs
  • Loading branch information
jwills committed May 3, 2023
1 parent 315d8be commit 8ac26f8
Showing 1 changed file with 18 additions and 3 deletions.
21 changes: 18 additions & 3 deletions dbt/adapters/duckdb/environments/local.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import threading

from . import Environment
from .. import credentials
from .. import utils
Expand All @@ -24,11 +26,13 @@ def execute(self, sql, bindings=None):


class DuckDBConnectionWrapper:
def __init__(self, cursor):
def __init__(self, cursor, env):
self._cursor = DuckDBCursorWrapper(cursor)
self._env = env

def close(self):
self._cursor.close()
self._env.notify_closed()

def cursor(self):
return self._cursor
Expand All @@ -39,14 +43,25 @@ def __init__(self, credentials: credentials.DuckDBCredentials):
# Set the conn attribute to None so it always exists even if
# DB initialization fails
self.conn = None
self.conn = self.initialize_db(credentials)
self._plugins = self.initialize_plugins(credentials)
self.creds = credentials
self.handle_count = 0
self.lock = threading.RLock()

def notify_closed(self):
with self.lock:
self.handle_count -= 1
if self.handle_count == 0 and self.creds.path != ":memory:":
self.close()

def handle(self):
# Extensions/settings need to be configured per cursor
with self.lock:
if self.conn is None:
self.conn = self.initialize_db(self.creds)
self.handle_count += 1
cursor = self.initialize_cursor(self.creds, self.conn.cursor())
return DuckDBConnectionWrapper(cursor)
return DuckDBConnectionWrapper(cursor, self)

def submit_python_job(self, handle, parsed_model: dict, compiled_code: str) -> AdapterResponse:
con = handle.cursor()
Expand Down

0 comments on commit 8ac26f8

Please sign in to comment.