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)