Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Have the local environment release the connection when it's not in use #163

Merged
merged 1 commit into from
May 3, 2023
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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