def execute_model(cls, profile, model):
parts = re.split(r'-- (DBT_OPERATION .*)', model.compiled_contents)
connection = cls.get_connection(profile)
if flags.STRICT_MODE:
validate_connection(connection)
for i, part in enumerate(parts):
matches = re.match(r'^DBT_OPERATION ({.*})$', part)
if matches is not None:
instruction_string = matches.groups()[0]
instruction = yaml.safe_load(instruction_string)
function = instruction['function']
kwargs = instruction['args']
def call_expand_target_column_types(kwargs):
kwargs.update({'profile': profile})
return cls.expand_target_column_types(**kwargs)
func_map = {
'expand_column_types_if_needed':
call_expand_target_column_types
}
func_map[function](kwargs)
else:
handle, cursor = cls.add_query_to_transaction(
part, connection, model.name)
handle.commit()
status = cls.get_status(cursor)
cursor.close()
return status
def close(cls, connection):
if dbt.flags.STRICT_MODE:
validate_connection(connection)
connection['state'] = 'closed'
return connection
def execute_model(cls,
profile,
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'))
validate_connection(connection)
model_name = model.get('name')
model_schema = model.get('schema')
dataset = cls.get_dataset(profile, model_schema, model_name)
if materialization == 'view':
res = cls.materialize_as_view(profile, dataset, model)
elif materialization == 'table':
res = cls.materialize_as_table(profile, dataset, model,
sql_override, decorator)
else:
msg = "Invalid relation type: '{}'".format(materialization)
raise dbt.exceptions.RuntimeException(msg, model)
return res
def query_for_existing(cls, profile, schema):
query = """
select TABLE_NAME as name, TABLE_TYPE as type
from INFORMATION_SCHEMA.TABLES
where TABLE_SCHEMA = '{schema}'
""".format(schema=schema).strip() # noqa
connection = cls.get_connection(profile)
if flags.STRICT_MODE:
validate_connection(connection)
_, cursor = cls.add_query_to_transaction(
query, connection, schema)
results = cursor.fetchall()
relation_type_lookup = {
'BASE TABLE': 'table',
'VIEW': 'view'
}
existing = [(name, relation_type_lookup.get(relation_type))
for (name, relation_type) in results]
return dict(existing)
def execute_model(cls, profile, model):
connection = cls.get_connection(profile, model.get('name'))
if flags.STRICT_MODE:
validate_connection(connection)
return super(PostgresAdapter, cls).execute_model(profile, model)
def get_columns_in_table(cls, profile, schema_name, table_name):
connection = cls.get_connection(profile)
if flags.STRICT_MODE:
validate_connection(connection)
query = """
select column_name, data_type, character_maximum_length
from information_schema.columns
where table_name = '{table_name}'
""".format(table_name=table_name).strip()
if schema_name is not None:
query += (" AND table_schema = '{schema_name}'".format(
schema_name=schema_name))
handle, cursor = cls.add_query_to_transaction(query, connection,
table_name)
data = cursor.fetchall()
columns = []
for row in data:
name, data_type, char_size = row
column = Column(name, data_type, char_size)
columns.append(column)
return columns
def expand_target_column_types(cls, profile, temp_table, to_schema,
to_table):
connection = cls.get_connection(profile)
if flags.STRICT_MODE:
validate_connection(connection)
reference_columns = {
col.name: col
for col in cls.get_columns_in_table(profile, None, temp_table)
}
target_columns = {
col.name: col
for col in cls.get_columns_in_table(profile, to_schema, to_table)
}
for column_name, reference_column in reference_columns.items():
target_column = target_columns.get(column_name)
if target_column is not None and \
target_column.can_expand_to(reference_column):
new_type = Column.string_type(reference_column.string_size())
logger.debug("Changing col type from %s to %s in table %s.%s",
target_column.data_type, new_type, to_schema,
to_table)
cls.alter_column_type(connection, to_schema, to_table,
column_name, new_type)
def get_default_schema(cls, profile):
connection = cls.get_connection(profile)
if flags.STRICT_MODE:
validate_connection(connection)
return connection.get('credentials', {}).get('schema')
def create_table(cls, profile, schema, table, columns, sort, dist):
connection = cls.get_connection(profile)
if flags.STRICT_MODE:
validate_connection(connection)
fields = [
'"{field}" {data_type}'.format(field=column.name,
data_type=column.data_type)
for column in columns
]
fields_csv = ",\n ".join(fields)
dist = cls.dist_qualifier(dist)
sort = cls.sort_qualifier('compound', sort)
sql = """
create table if not exists "{schema}"."{table}" (
{fields}
)
{dist} {sort}
""".format(schema=schema,
table=table,
fields=fields_csv,
sort=sort,
dist=dist)
logger.debug('creating table "%s"."%s"'.format(schema, table))
cls.add_query_to_transaction(sql, connection, table)
def commit(cls, profile):
connection = cls.get_connection(profile)
if flags.STRICT_MODE:
validate_connection(connection)
handle = connection.get('handle')
handle.commit()
def execute_one(cls, profile, query, model_name=None):
connection = cls.get_connection(profile)
if flags.STRICT_MODE:
validate_connection(connection)
handle = connection.get('handle')
return cls.add_query_to_transaction(query, connection, model_name)
def create_schema(cls, profile, schema, model_name=None):
connection = cls.get_connection(profile)
if flags.STRICT_MODE:
validate_connection(connection)
query = ('create schema if not exists "{schema}"'.format(
schema=schema))
handle, cursor = cls.add_query_to_transaction(query, connection,
model_name)
def rename(cls, profile, from_name, to_name, model_name=None):
connection = cls.get_connection(profile)
if flags.STRICT_MODE:
validate_connection(connection)
schema = connection.get('credentials', {}).get('schema')
query = ('alter table "{schema}"."{from_name}" rename to "{to_name}"'.
format(schema=schema, from_name=from_name, to_name=to_name))
handle, cursor = cls.add_query_to_transaction(query, connection,
model_name)
def truncate(cls, profile, table, model_name=None):
connection = cls.get_connection(profile)
if flags.STRICT_MODE:
validate_connection(connection)
schema = connection.get('credentials', {}).get('schema')
query = ('truncate table "{schema}"."{table}"'.format(schema=schema,
table=table))
handle, cursor = cls.add_query_to_transaction(query, connection,
model_name)
def drop_table(cls, profile, table, model_name):
connection = cls.get_connection(profile)
if flags.STRICT_MODE:
validate_connection(connection)
schema = connection.get('credentials', {}).get('schema')
query = ('drop table if exists "{schema}"."{table}" cascade'.format(
schema=schema, table=table))
handle, cursor = cls.add_query_to_transaction(query, connection,
model_name)
def close(cls, connection):
if dbt.flags.STRICT_MODE:
validate_connection(connection)
connection = cls.reload(connection)
if connection.get('state') == 'closed':
return connection
connection.get('handle').close()
connection['state'] = 'closed'
connections_in_use[connection.get('name')] = connection
return connection
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:
validate_connection(result)
return cls.open_connection(result)
finally:
lock.release()
def execute_all(cls, profile, queries, model_name=None):
if len(queries) == 0:
return
connection = cls.get_connection(profile)
if flags.STRICT_MODE:
validate_connection(connection)
handle = connection.get('handle')
for i, query in enumerate(queries):
handle, cursor = cls.add_query_to_transaction(
query, connection, model_name)
return cls.get_status(cursor)
def begin(cls, profile, name='master'):
global connections_in_use
connection = cls.get_connection(profile, name)
if dbt.flags.STRICT_MODE:
validate_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 rollback(cls, connection):
if dbt.flags.STRICT_MODE:
validate_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 execute_model(cls, profile, model):
parts = re.split(r'-- (DBT_OPERATION .*)', model.compiled_contents)
connection = cls.get_connection(profile)
if flags.STRICT_MODE:
validate_connection(connection)
# snowflake requires a schema to be specified for temporary tables
# TODO setup templates to be adapter-specific. then we can just use
# the existing schema for temp tables.
cls.add_query_to_transaction(
'USE SCHEMA "{}"'.format(
connection.get('credentials', {}).get('schema')),
connection)
for i, part in enumerate(parts):
matches = re.match(r'^DBT_OPERATION ({.*})$', part)
if matches is not None:
instruction_string = matches.groups()[0]
instruction = yaml.safe_load(instruction_string)
function = instruction['function']
kwargs = instruction['args']
def call_expand_target_column_types(kwargs):
kwargs.update({'profile': profile})
return cls.expand_target_column_types(**kwargs)
func_map = {
'expand_column_types_if_needed':
call_expand_target_column_types
}
func_map[function](kwargs)
else:
handle, cursor = cls.add_query_to_transaction(
part, connection, model.name)
handle.commit()
status = cls.get_status(cursor)
cursor.close()
return status
def query_for_existing(cls, profile, schema):
query = """
select tablename as name, 'table' as type from pg_tables
where schemaname = '{schema}'
union all
select viewname as name, 'view' as type from pg_views
where schemaname = '{schema}'
""".format(schema=schema).strip() # noqa
connection = cls.get_connection(profile)
if flags.STRICT_MODE:
validate_connection(connection)
_, cursor = cls.add_query_to_transaction(query, connection, schema)
results = cursor.fetchall()
existing = [(name, relation_type) for (name, relation_type) in results]
return dict(existing)
def commit(cls, profile, connection):
global connections_in_use
if dbt.flags.STRICT_MODE:
validate_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 rename(cls, profile, from_name, to_name, model_name=None):
connection = cls.get_connection(profile)
if flags.STRICT_MODE:
validate_connection(connection)
schema = connection.get('credentials', {}).get('schema')
# in snowflake, if you fail to include the quoted schema in the
# identifier, the new table will have `schema.upper()` as its new
# schema
query = ('''
alter table "{schema}"."{from_name}"
rename to "{schema}"."{to_name}"
'''.format(
schema=schema,
from_name=from_name,
to_name=to_name)).strip()
handle, cursor = cls.add_query_to_transaction(
query, connection, model_name)
def acquire_connection(cls, profile):
# profile requires some marshalling right now because it includes a
# wee bit of global config.
# TODO remove this
credentials = copy.deepcopy(profile)
credentials.pop('type', None)
credentials.pop('threads', None)
result = {
'type': 'postgres',
'state': 'init',
'handle': None,
'credentials': credentials
}
logger.info('Connecting to postgres.')
if flags.STRICT_MODE:
validate_connection(result)
return cls.open_connection(result)
