def open(cls, connection: Connection) -> Connection:
if connection.state == "open":
logger.debug("Connection is already open, skipping open.")
return connection
try:
creds: AthenaCredentials = connection.credentials
handle = AthenaConnection(s3_staging_dir=creds.s3_staging_dir,
region_name=creds.region_name,
schema_name=creds.schema,
work_group=creds.work_group,
cursor_class=AthenaCursor,
formatter=AthenaParameterFormatter(),
poll_interval=creds.poll_interval,
profile_name=creds.aws_profile_name)
connection.state = "open"
connection.handle = handle
except Exception as e:
logger.debug("Got an error when attempting to open a Athena "
"connection: '{}'".format(e))
connection.handle = None
connection.state = "fail"
raise FailedToConnectException(str(e))
return connection
def open(cls, connection: Connection) -> Connection:
if connection.state == 'open':
logger.debug('Connection is already open, skipping open.')
return connection
credentials = connection.credentials
try:
connection_string = '{}:{}/{}'
handle = oracle.connect(
credentials.username,
credentials.password,
connection_string.format(credentials.host, credentials.port, credentials.database),
mode=oracle.SYSDBA if credentials.as_sysdba else oracle.DEFAULT_AUTH
)
connection.state = 'open'
connection.handle = handle
if credentials.nls_date_format is not None:
handle.cursor().execute(f"ALTER SESSION SET NLS_DATE_FORMAT = '{credentials.nls_date_format}'")
handle.cursor().execute(f"ALTER SESSION SET NLS_TIMESTAMP_FORMAT = '{credentials.nls_date_format}XFF'")
handle.cursor().execute(f"ALTER SESSION SET NLS_TIMESTAMP_TZ_FORMAT = '{credentials.nls_date_format}XFF TZR'")
except Exception as e:
logger.debug(
f"Got an error when attempting to open a oracle connection. Connection string = {credentials.connection_string}. Error: {e}"
)
connection.handle = None
connection.state = 'fail'
raise FailedToConnectException(f"{str(e)} using connection string {credentials.connection_string}")
return connection
def execute_model(cls, profile, project_cfg, model,
materialization, sql_override=None,
decorator=None, model_name=None):
if sql_override is None:
sql_override = model.get('injected_sql')
if flags.STRICT_MODE:
connection = cls.get_connection(profile, model.get('name'))
Connection(**connection)
model_name = model.get('name')
model_schema = model.get('schema')
dataset = cls.get_dataset(profile, project_cfg,
model_schema, model_name)
if materialization == 'view':
res = cls.materialize_as_view(profile, project_cfg, dataset, model)
elif materialization == 'table':
res = cls.materialize_as_table(
profile, project_cfg, dataset, model,
sql_override, decorator)
else:
msg = "Invalid relation type: '{}'".format(materialization)
raise dbt.exceptions.RuntimeException(msg, model)
return res
def close(cls, connection):
if dbt.flags.STRICT_MODE:
Connection(**connection)
connection['state'] = 'closed'
return connection
def acquire_connection(self, name):
global connections_available, lock
# we add a magic number, 2 because there are overhead connections,
# one for pre- and post-run hooks and other misc operations that occur
# before the run starts, and one for integration tests.
max_connections = self.config.threads + 2
with lock:
num_allocated = self.total_connections_allocated()
if len(connections_available) > 0:
logger.debug('Re-using an available connection from the pool.')
to_return = connections_available.pop()
to_return.name = name
return to_return
elif num_allocated >= max_connections:
raise dbt.exceptions.InternalException(
'Tried to request a new connection "{}" but '
'the maximum number of connections are already '
'allocated!'.format(name))
logger.debug(
'Opening a new connection ({} currently allocated)'.format(
num_allocated))
result = Connection(type=self.type(),
name=name,
state='init',
transaction_open=False,
handle=None,
credentials=self.config.credentials)
return self.open_connection(result)
def _empty_connection(self):
return Connection(type=self.TYPE,
name=None,
state='init',
transaction_open=False,
handle=None,
credentials=self.profile.credentials)
def close(cls, connection: Connection) -> Connection:
if dbt.flags.STRICT_MODE:
if not isinstance(connection, Connection):
raise dbt.exceptions.CompilerException(
f'In close, got {connection} - not a Connection!')
# if the connection is in closed or init, there's nothing to do
if connection.state in {ConnectionState.CLOSED, ConnectionState.INIT}:
return connection
if connection.transaction_open and connection.handle:
cls._rollback_handle(connection)
connection.transaction_open = False
cls._close_handle(connection)
connection.state = ConnectionState.CLOSED
return connection
def set_connection_name(self, name: Optional[str] = None) -> Connection:
conn_name: str
if name is None:
# if a name isn't specified, we'll re-use a single handle
# named 'master'
conn_name = 'master'
else:
if not isinstance(name, str):
raise dbt.exceptions.CompilerException(
f'For connection name, got {name} - not a string!'
)
assert isinstance(name, str)
conn_name = name
conn = self.get_if_exists()
if conn is None:
conn = Connection(
type=Identifier(self.TYPE),
name=None,
state=ConnectionState.INIT,
transaction_open=False,
handle=None,
credentials=self.profile.credentials
)
self.set_thread_connection(conn)
if conn.name == conn_name and conn.state == 'open':
return conn
logger.debug(
'Acquiring new {} connection "{}".'.format(self.TYPE, conn_name))
if conn.state == 'open':
logger.debug(
'Re-using an available connection from the pool (formerly {}).'
.format(conn.name)
)
else:
conn.handle = LazyHandle(self.open)
conn.name = conn_name
return conn
def open(cls, connection: Connection) -> Connection:
if connection.state == "open":
logger.debug("Connection is already open, skipping open")
return connection
credentials = connection.credentials
kwargs = {
"host": credentials.host,
"port": credentials.port,
"user": credentials.username,
"password": credentials.password,
}
try:
connection.handle = MySQLdb.connect(**kwargs)
connection.state = ConnectionState.OPEN
except MySQLdb.Error as e:
logger.debug(f"Error connecting to database: {e}")
connection.handle = None
connection.state = ConnectionState.FAIL
raise exceptions.FailedToConnectException(str(e))
return connection
def open(cls, connection: Connection):
if connection.state == ConnectionState.OPEN:
logger.debug("Connection is already open, skipping open.")
return connection
try:
connection.handle = pyodbc.connect(
connection.credentials.connection_string(),
autocommit=True,
)
connection.state = ConnectionState.OPEN
except pyodbc.OperationalError as e:
logger.debug(
f"Got an error when attempting to open an odbc connection: '{e}'"
)
connection.handle = None
connection.state = ConnectionState.FAIL
raise dbt.exceptions.FailedToConnectException(str(e)) from e
return connection
def set_connection_name(self, name=None):
if name is None:
# if a name isn't specified, we'll re-use a single handle
# named 'master'
name = 'master'
conn = self.get_if_exists()
thread_id_key = self.get_thread_identifier()
if conn is None:
conn = Connection(type=self.TYPE,
name=None,
state='init',
transaction_open=False,
handle=None,
credentials=self.profile.credentials)
self.thread_connections[thread_id_key] = conn
if conn.name == name and conn.state == 'open':
return conn
logger.debug('Acquiring new {} connection "{}".'.format(
self.TYPE, name))
if conn.state == 'open':
logger.debug(
'Re-using an available connection from the pool (formerly {}).'
.format(conn.name))
else:
logger.debug(
'Opening a new connection, currently in state {}'.format(
conn.state))
self.open(conn)
conn.name = name
return conn
def acquire_connection(cls, profile, name):
global connections_available, lock
# we add a magic number, 2 because there are overhead connections,
# one for pre- and post-run hooks and other misc operations that occur
# before the run starts, and one for integration tests.
max_connections = profile.get('threads', 1) + 2
try:
lock.acquire()
num_allocated = cls.total_connections_allocated()
if len(connections_available) > 0:
logger.debug('Re-using an available connection from the pool.')
to_return = connections_available.pop()
to_return['name'] = name
return to_return
elif num_allocated >= max_connections:
raise dbt.exceptions.InternalException(
'Tried to request a new connection "{}" but '
'the maximum number of connections are already '
'allocated!'.format(name))
logger.debug(
'Opening a new connection ({} currently allocated)'.format(
num_allocated))
credentials = copy.deepcopy(profile)
credentials.pop('type', None)
credentials.pop('threads', None)
result = {
'type': cls.type(),
'name': name,
'state': 'init',
'transaction_open': False,
'handle': None,
'credentials': credentials
}
if dbt.flags.STRICT_MODE:
Connection(**result)
return cls.open_connection(result)
finally:
lock.release()
def _rollback(cls, connection: Connection) -> None:
"""Roll back the given connection."""
if dbt.flags.STRICT_MODE:
if not isinstance(connection, Connection):
raise dbt.exceptions.CompilerException(
f'In _rollback, got {connection} - not a Connection!')
if connection.transaction_open is False:
raise dbt.exceptions.InternalException(
'Tried to rollback transaction on connection "{}", but '
'it does not have one open!'.format(connection.name))
logger.debug('On {}: ROLLBACK'.format(connection.name))
cls._rollback_handle(connection)
connection.transaction_open = False
def rollback(self, connection):
if dbt.flags.STRICT_MODE:
Connection(**connection)
connection = self.reload(connection)
if connection.transaction_open is False:
raise dbt.exceptions.InternalException(
'Tried to rollback transaction on connection "{}", but '
'it does not have one open!'.format(connection.name))
logger.debug('On {}: ROLLBACK'.format(connection.name))
connection.handle.rollback()
connection.transaction_open = False
connections_in_use[connection.name] = connection
return connection
def rollback(cls, connection):
if dbt.flags.STRICT_MODE:
Connection(**connection)
connection = cls.reload(connection)
if connection['transaction_open'] is False:
raise dbt.exceptions.InternalException(
'Tried to rollback transaction on connection "{}", but '
'it does not have one open!'.format(connection.get('name')))
logger.debug('On {}: ROLLBACK'.format(connection.get('name')))
connection.get('handle').rollback()
connection['transaction_open'] = False
connections_in_use[connection.get('name')] = connection
return connection
def begin(cls, profile, name='master'):
global connections_in_use
connection = cls.get_connection(profile, name)
if dbt.flags.STRICT_MODE:
Connection(**connection)
if connection['transaction_open'] is True:
raise dbt.exceptions.InternalException(
'Tried to begin a new transaction on connection "{}", but '
'it already had one open!'.format(connection.get('name')))
cls.add_begin_query(profile, name)
connection['transaction_open'] = True
connections_in_use[name] = connection
return connection
def commit(cls, profile, connection):
global connections_in_use
if dbt.flags.STRICT_MODE:
Connection(**connection)
connection = cls.reload(connection)
if connection['transaction_open'] is False:
raise dbt.exceptions.InternalException(
'Tried to commit transaction on connection "{}", but '
'it does not have one open!'.format(connection.get('name')))
logger.debug('On {}: COMMIT'.format(connection.get('name')))
cls.add_commit_query(profile, connection.get('name'))
connection['transaction_open'] = False
connections_in_use[connection.get('name')] = connection
return connection
def open(cls, connection: Connection):
if connection.state == "open":
logger.debug("Connection is already open, skipping open.")
return connection
credentials: SQLiteCredentials = connection.credentials
schemas_and_paths = {}
for schema, path in credentials.schemas_and_paths.items():
# Make .db file path absolute
schemas_and_paths[schema] = os.path.abspath(path)
try:
if 'main' in schemas_and_paths:
handle: sqlite3.Connection = sqlite3.connect(
schemas_and_paths['main'])
else:
raise FailedToConnectException(
"at least one schema must be called 'main'")
if len(credentials.extensions) > 0:
handle.enable_load_extension(True)
for ext_path in credentials.extensions:
handle.load_extension(ext_path)
cursor = handle.cursor()
attached = []
for schema in set(schemas_and_paths.keys()) - set(['main']):
path = schemas_and_paths[schema]
cursor.execute(f"attach '{path}' as '{schema}'")
attached.append(schema)
for path in glob.glob(
os.path.join(credentials.schema_directory, "*.db")):
abs_path = os.path.abspath(path)
# if file was already attached from being defined in schemas_and_paths, ignore it
if not abs_path in schemas_and_paths.values():
schema = os.path.basename(path)[:-3]
# has schema name been used already?
if schema not in attached:
cursor.execute(f"attach '{path}' as '{schema}'")
else:
raise FailedToConnectException(
f"found {path} while scanning schema_directory, but cannot attach it as '{schema}' "
+
f"because that schema name is already defined in schemas_and_paths. "
+ f"fix your ~/.dbt/profiles.yml file")
# # uncomment these lines to print out SQL: this only happens if statement is successful
# handle.set_trace_callback(print)
# sqlite3.enable_callback_tracebacks(True)
connection.state = "open"
connection.handle = handle
return connection
except sqlite3.Error as e:
logger.debug(
"Got an error when attempting to open a sqlite3 connection: '%s'",
e)
connection.handle = None
connection.state = "fail"
raise FailedToConnectException(str(e))
except Exception as e:
print(f"Unknown error opening SQLite connection: {e}")
raise e
