def get_latest_materialization_events( self, asset_keys: Sequence[AssetKey] ) -> Mapping[AssetKey, Optional[EventLogEntry]]: check.list_param(asset_keys, "asset_keys", AssetKey) rows = self._fetch_asset_rows(asset_keys=asset_keys) to_backcompat_fetch = set() results: Dict[AssetKey, Optional[EventLogEntry]] = {} for row in rows: asset_key = AssetKey.from_db_string(row[0]) if not asset_key: continue event_or_materialization = ( deserialize_json_to_dagster_namedtuple(row[1]) if row[1] else None ) if isinstance(event_or_materialization, EventLogEntry): results[asset_key] = event_or_materialization else: to_backcompat_fetch.add(asset_key) if to_backcompat_fetch: latest_event_subquery = ( db.select( [ SqlEventLogStorageTable.c.asset_key, db.func.max(SqlEventLogStorageTable.c.timestamp).label("timestamp"), ] ) .where( db.and_( SqlEventLogStorageTable.c.asset_key.in_( [asset_key.to_string() for asset_key in to_backcompat_fetch] ), SqlEventLogStorageTable.c.dagster_event_type == DagsterEventType.ASSET_MATERIALIZATION.value, ) ) .group_by(SqlEventLogStorageTable.c.asset_key) .subquery() ) backcompat_query = db.select( [SqlEventLogStorageTable.c.asset_key, SqlEventLogStorageTable.c.event] ).join( latest_event_subquery, db.and_( SqlEventLogStorageTable.c.asset_key == latest_event_subquery.c.asset_key, SqlEventLogStorageTable.c.timestamp == latest_event_subquery.c.timestamp, ), ) with self.index_connection() as conn: event_rows = conn.execute(backcompat_query).fetchall() for row in event_rows: asset_key = AssetKey.from_db_string(row[0]) if asset_key: results[asset_key] = cast( EventLogEntry, deserialize_json_to_dagster_namedtuple(row[1]) ) return results
def all_asset_keys(self): with self.index_connection() as conn: results = conn.execute( db.select([AssetKeyTable.c.asset_key, AssetKeyTable.c.asset_details]) ).fetchall() asset_keys = set() wiped = set() wiped_timestamps = {} for result in results: asset_key = AssetKey.from_db_string(result[0]) asset_details: Optional[AssetDetails] = AssetDetails.from_db_string(result[1]) asset_keys.add(asset_key) if asset_details and asset_details.last_wipe_timestamp: wiped_timestamps[asset_key] = asset_details.last_wipe_timestamp if wiped_timestamps: materialized_timestamps = {} # fetch the last materialization timestamp per asset key materialization_results = conn.execute( db.select( [ SqlEventLogStorageTable.c.asset_key, db.func.max(SqlEventLogStorageTable.c.timestamp), ] ) .where( SqlEventLogStorageTable.c.asset_key.in_( [asset_key.to_string() for asset_key in wiped_timestamps.keys()] ) ) .group_by(SqlEventLogStorageTable.c.asset_key) .order_by(db.func.max(SqlEventLogStorageTable.c.timestamp).asc()) ).fetchall() for result in materialization_results: asset_key = AssetKey.from_db_string(result[0]) last_materialized_timestamp = result[1] materialized_timestamps[asset_key] = last_materialized_timestamp # calculate the set of wiped asset keys that have not had a materialization since # the wipe timestamp wiped = set( [ asset_key for asset_key in wiped_timestamps.keys() if not materialized_timestamps.get(asset_key) or utc_datetime_from_naive(materialized_timestamps.get(asset_key)) < utc_datetime_from_timestamp(wiped_timestamps[asset_key]) ] ) return list(asset_keys.difference(wiped))
def get_all_asset_keys(self, prefix_path=None): if not prefix_path: if self.has_secondary_index(SECONDARY_INDEX_ASSET_KEY): query = db.select([AssetKeyTable.c.asset_key]) else: query = (db.select([ SqlEventLogStorageTable.c.asset_key ]).where( SqlEventLogStorageTable.c.asset_key != None).distinct()) else: if self.has_secondary_index(SECONDARY_INDEX_ASSET_KEY): query = db.select([AssetKeyTable.c.asset_key]).where( AssetKeyTable.c.asset_key.startswith( AssetKey.get_db_prefix(prefix_path))) else: query = (db.select([ SqlEventLogStorageTable.c.asset_key ]).where(SqlEventLogStorageTable.c.asset_key != None).where( SqlEventLogStorageTable.c.asset_key.startswith( AssetKey.get_db_prefix(prefix_path))).distinct()) with self.connect() as conn: results = conn.execute(query).fetchall() return [ AssetKey.from_db_string(asset_key) for (asset_key, ) in results if asset_key ]
def asset_wipe_command(key, **cli_args): if not cli_args.get("all") and len(key) == 0: raise click.UsageError( "Error, you must specify an asset key or use `--all` to wipe all asset keys." ) if cli_args.get("all") and len(key) > 0: raise click.UsageError("Error, cannot use more than one of: asset key, `--all`.") with DagsterInstance.get() as instance: if len(key) > 0: asset_keys = [AssetKey.from_db_string(key_string) for key_string in key] prompt = ( "Are you sure you want to remove the asset key indexes for these keys from the event " "logs? Type DELETE" ) else: asset_keys = instance.all_asset_keys() prompt = "Are you sure you want to remove all asset key indexes from the event logs? Type DELETE" confirmation = click.prompt(prompt) if confirmation == "DELETE": with DagsterInstance.get() as instance: instance.wipe_assets(asset_keys) click.echo("Removed asset indexes from event logs") else: click.echo("Exiting without removing asset indexes")
def asset_wipe_command(key, **cli_args): if not cli_args.get('all') and len(key) == 0: raise click.UsageError( 'Error, you must specify an asset key or use `--all` to wipe all asset keys.' ) if cli_args.get('all') and len(key) > 0: raise click.UsageError( 'Error, cannot use more than one of: asset key, `--all`.') if len(key) > 0: asset_keys = [ AssetKey.from_db_string(key_string) for key_string in key ] prompt = ( 'Are you sure you want to remove the asset key indexes for these keys from the event ' 'logs? Type DELETE') else: asset_keys = None prompt = ( 'Are you sure you want to remove all asset key indexes from the event logs? Type DELETE' ) confirmation = click.prompt(prompt) if confirmation == 'DELETE': instance = DagsterInstance.get() if asset_keys: instance.wipe_assets(asset_keys) else: instance.wipe_all_assets() click.echo('Removed asset indexes from event logs') else: click.echo('Exiting without removing asset indexes')
def get_asset_keys( self, prefix: Optional[List[str]] = None, limit: Optional[int] = None, cursor: Optional[str] = None, ) -> Iterable[AssetKey]: rows = self._fetch_asset_rows(prefix=prefix, limit=limit, cursor=cursor) asset_keys = [AssetKey.from_db_string(row[0]) for row in sorted(rows, key=lambda x: x[0])] return [asset_key for asset_key in asset_keys if asset_key]
def get_all_asset_keys(self): query = db.select([SqlEventLogStorageTable.c.asset_key]).distinct() with self.connect() as conn: results = conn.execute(query).fetchall() return [ AssetKey.from_db_string(asset_key) for (asset_key, ) in results if asset_key ]
def get_asset_keys(self, prefix_path=None): lazy_migrate = False if not prefix_path: if self.has_secondary_index(SECONDARY_INDEX_ASSET_KEY): query = db.select([AssetKeyTable.c.asset_key]) else: query = (db.select([ SqlEventLogStorageTable.c.asset_key ]).where( SqlEventLogStorageTable.c.asset_key != None).distinct()) # This is in place to migrate everyone to using the secondary index table for asset # keys. Performing this migration should result in a big performance boost for # any asset-catalog reads. # After a sufficient amount of time (>= 0.11.0?), we can remove the checks # for has_secondary_index(SECONDARY_INDEX_ASSET_KEY) and always read from the # AssetKeyTable, since we are already writing to the table. Tracking the conditional # check removal here: https://github.com/dagster-io/dagster/issues/3507 lazy_migrate = True else: if self.has_secondary_index(SECONDARY_INDEX_ASSET_KEY): query = db.select([AssetKeyTable.c.asset_key]).where( db.or_( AssetKeyTable.c.asset_key.startswith( AssetKey.get_db_prefix(prefix_path)), AssetKeyTable.c.asset_key.startswith( AssetKey.get_db_prefix(prefix_path, legacy=True)), )) else: query = (db.select([ SqlEventLogStorageTable.c.asset_key ]).where(SqlEventLogStorageTable.c.asset_key != None).where( db.or_( SqlEventLogStorageTable.c.asset_key.startswith( AssetKey.get_db_prefix(prefix_path)), SqlEventLogStorageTable.c.asset_key.startswith( AssetKey.get_db_prefix(prefix_path, legacy=True)), )).distinct()) with self.index_connection() as conn: results = conn.execute(query).fetchall() if lazy_migrate: # This is in place to migrate everyone to using the secondary index table for asset # keys. Performing this migration should result in a big performance boost for # any subsequent asset-catalog reads. self._lazy_migrate_secondary_index_asset_key( [asset_key for (asset_key, ) in results if asset_key]) return list( set([ AssetKey.from_db_string(asset_key) for (asset_key, ) in results if asset_key ]))
def delete_events_for_run(self, conn, run_id): check.str_param(run_id, "run_id") delete_statement = ( SqlEventLogStorageTable.delete().where( # pylint: disable=no-value-for-parameter SqlEventLogStorageTable.c.run_id == run_id ) ) removed_asset_key_query = ( db.select([SqlEventLogStorageTable.c.asset_key]) .where(SqlEventLogStorageTable.c.run_id == run_id) .where(SqlEventLogStorageTable.c.asset_key != None) .group_by(SqlEventLogStorageTable.c.asset_key) ) removed_asset_keys = [ AssetKey.from_db_string(row[0]) for row in conn.execute(removed_asset_key_query).fetchall() ] conn.execute(delete_statement) if len(removed_asset_keys) > 0: keys_to_check = [] keys_to_check.extend([key.to_string() for key in removed_asset_keys]) keys_to_check.extend([key.to_string(legacy=True) for key in removed_asset_keys]) remaining_asset_keys = [ AssetKey.from_db_string(row[0]) for row in conn.execute( db.select([SqlEventLogStorageTable.c.asset_key]) .where(SqlEventLogStorageTable.c.asset_key.in_(keys_to_check)) .group_by(SqlEventLogStorageTable.c.asset_key) ) ] to_remove = set(removed_asset_keys) - set(remaining_asset_keys) if to_remove: keys_to_remove = [] keys_to_remove.extend([key.to_string() for key in to_remove]) keys_to_remove.extend([key.to_string(legacy=True) for key in to_remove]) conn.execute( AssetKeyTable.delete().where( # pylint: disable=no-value-for-parameter AssetKeyTable.c.asset_key.in_(keys_to_remove) ) )
def all_asset_tags(self): query = db.select([AssetKeyTable.c.asset_key, AssetKeyTable.c.last_materialization]) tags_by_asset_key = defaultdict(dict) with self.index_connection() as conn: rows = conn.execute(query).fetchall() for asset_key, json_str in rows: materialization = self._asset_materialization_from_json_column(json_str) if materialization and materialization.tags: tags_by_asset_key[AssetKey.from_db_string(asset_key)] = { k: v for k, v in materialization.tags.items() } return tags_by_asset_key
def all_asset_tags(self): query = db.select([AssetKeyTable.c.asset_key, AssetKeyTable.c.last_materialization]) tags_by_asset_key = defaultdict(dict) with self.index_connection() as conn: rows = conn.execute(query).fetchall() for asset_key, materialization_str in rows: if materialization_str: materialization = deserialize_json_to_dagster_namedtuple(materialization_str) tags_by_asset_key[AssetKey.from_db_string(asset_key)] = { k: v for k, v in (materialization.tags or {}).items() } return tags_by_asset_key
def _lazy_migrate_secondary_index_asset_key(self, conn, asset_keys): results = conn.execute(db.select([AssetKeyTable.c.asset_key ])).fetchall() existing = [asset_key for (asset_key, ) in results if asset_key] to_migrate = set(asset_keys) - set(existing) for asset_key in to_migrate: try: conn.execute(AssetKeyTable.insert().values( # pylint: disable=no-value-for-parameter asset_key=AssetKey.from_db_string(asset_key).to_string())) except db.exc.IntegrityError: # asset key already present pass self.enable_secondary_index(SECONDARY_INDEX_ASSET_KEY)
def get_all_asset_keys(self, prefix_path=None): if not prefix_path: query = db.select([SqlEventLogStorageTable.c.asset_key]).distinct() else: query = (db.select([SqlEventLogStorageTable.c.asset_key]).where( SqlEventLogStorageTable.c.asset_key.startswith( AssetKey.get_db_prefix(prefix_path))).distinct()) with self.connect() as conn: results = conn.execute(query).fetchall() return [ AssetKey.from_db_string(asset_key) for (asset_key, ) in results if asset_key ]
def all_asset_tags(self): tags_by_asset_key = defaultdict(dict) if self.has_secondary_index(ASSET_KEY_INDEX_COLS): query = (db.select([ AssetKeyTable.c.asset_key, AssetKeyTable.c.tags ]).where(AssetKeyTable.c.tags != None).where( db.or_( AssetKeyTable.c.wipe_timestamp == None, AssetKeyTable.c.last_materialization_timestamp > AssetKeyTable.c.wipe_timestamp, ))) with self.index_connection() as conn: rows = conn.execute(query).fetchall() for asset_key, tags_json in rows: tags = seven.json.loads(tags_json) if tags: tags_by_asset_key[AssetKey.from_db_string( asset_key)] = tags else: query = db.select([ AssetKeyTable.c.asset_key, AssetKeyTable.c.last_materialization ]) with self.index_connection() as conn: rows = conn.execute(query).fetchall() for asset_key, json_str in rows: materialization = self._asset_materialization_from_json_column( json_str) if materialization and materialization.tags: tags_by_asset_key[AssetKey.from_db_string( asset_key)] = { k: v for k, v in materialization.tags.items() } return tags_by_asset_key
def get_materialization_count_by_partition( self, asset_keys: Sequence[AssetKey] ) -> Mapping[AssetKey, Mapping[str, int]]: check.list_param(asset_keys, "asset_keys", AssetKey) query = ( db.select( [ SqlEventLogStorageTable.c.asset_key, SqlEventLogStorageTable.c.partition, db.func.count(SqlEventLogStorageTable.c.id), ] ) .where( db.and_( db.or_( SqlEventLogStorageTable.c.asset_key.in_( [asset_key.to_string() for asset_key in asset_keys] ), SqlEventLogStorageTable.c.asset_key.in_( [asset_key.to_string(legacy=True) for asset_key in asset_keys] ), ), SqlEventLogStorageTable.c.partition != None, ) ) .group_by(SqlEventLogStorageTable.c.asset_key, SqlEventLogStorageTable.c.partition) ) assets_details = self._get_assets_details(asset_keys) query = self._add_assets_wipe_filter_to_query(query, assets_details, asset_keys) with self.index_connection() as conn: results = conn.execute(query).fetchall() materialization_count_by_partition: Dict[AssetKey, Dict[str, int]] = { asset_key: {} for asset_key in asset_keys } for row in results: asset_key = AssetKey.from_db_string(row[0]) if asset_key: materialization_count_by_partition[asset_key][row[1]] = row[2] return materialization_count_by_partition
def get_asset_keys( self, prefix: Optional[List[str]] = None, limit: Optional[int] = None, cursor: Optional[str] = None, ) -> Iterable[AssetKey]: # base implementation of get_asset_keys, using the existing `all_asset_keys` and doing the # filtering in-memory asset_keys = sorted(self.all_asset_keys(), key=str) if prefix: asset_keys = [ asset_key for asset_key in asset_keys if asset_key.path[: len(prefix)] == prefix ] if cursor: cursor_asset = AssetKey.from_db_string(cursor) if cursor_asset and cursor_asset in asset_keys: idx = asset_keys.index(cursor_asset) asset_keys = asset_keys[idx + 1 :] if limit: asset_keys = asset_keys[:limit] return asset_keys
def _fetch_backcompat_materialization_times(self, asset_keys): # fetches the latest materialization timestamp for the given asset_keys. Uses the (slower) # raw event log table. backcompat_query = ( db.select( [ SqlEventLogStorageTable.c.asset_key, db.func.max(SqlEventLogStorageTable.c.timestamp), ] ) .where( SqlEventLogStorageTable.c.asset_key.in_( [asset_key.to_string() for asset_key in asset_keys] ) ) .group_by(SqlEventLogStorageTable.c.asset_key) .order_by(db.func.max(SqlEventLogStorageTable.c.timestamp).asc()) ) with self.index_connection() as conn: backcompat_rows = conn.execute(backcompat_query).fetchall() return {AssetKey.from_db_string(row[0]): row[1] for row in backcompat_rows}
def _fetch_raw_asset_rows(self, asset_keys=None, prefix=None, limit=None, cursor=None): # fetches rows containing asset_key, last_materialization, and asset_details from the DB, # applying the filters specified in the arguments. Does not guarantee that the number of # rows returned will match the limit specified. This helper function is used to fetch a # chunk of asset key rows, which may or may not be wiped. # # Returns a tuple of (rows, has_more, cursor), where each row is a tuple of serialized # asset_key, materialization, and asset_details columns = [ AssetKeyTable.c.asset_key, AssetKeyTable.c.last_materialization, AssetKeyTable.c.asset_details, ] is_partial_query = bool(asset_keys) or bool(prefix) or bool(limit) or bool(cursor) if self.has_asset_key_index_cols() and not is_partial_query: # if the schema has been migrated, fetch the last_materialization_timestamp to see if # we can lazily migrate the data table columns.append(AssetKeyTable.c.last_materialization_timestamp) columns.append(AssetKeyTable.c.wipe_timestamp) query = db.select(columns).order_by(AssetKeyTable.c.asset_key.asc()) query = self._apply_asset_filter_to_query(query, asset_keys, prefix, limit, cursor) if self.has_secondary_index(ASSET_KEY_INDEX_COLS): query = query.where( db.or_( AssetKeyTable.c.wipe_timestamp == None, AssetKeyTable.c.last_materialization_timestamp > AssetKeyTable.c.wipe_timestamp, ) ) with self.index_connection() as conn: rows = conn.execute(query).fetchall() return rows, False, None with self.index_connection() as conn: rows = conn.execute(query).fetchall() wiped_timestamps_by_asset_key = {} row_by_asset_key = OrderedDict() for row in rows: asset_key = AssetKey.from_db_string(row[0]) if not asset_key: continue asset_details = AssetDetails.from_db_string(row[2]) if not asset_details or not asset_details.last_wipe_timestamp: row_by_asset_key[asset_key] = row continue materialization_or_event = ( deserialize_json_to_dagster_namedtuple(row[1]) if row[1] else None ) if isinstance(materialization_or_event, EventLogEntry): if asset_details.last_wipe_timestamp > materialization_or_event.timestamp: # this asset has not been materialized since being wiped, skip continue else: # add the key row_by_asset_key[asset_key] = row else: row_by_asset_key[asset_key] = row wiped_timestamps_by_asset_key[asset_key] = asset_details.last_wipe_timestamp if wiped_timestamps_by_asset_key: materialization_times = self._fetch_backcompat_materialization_times( wiped_timestamps_by_asset_key.keys() ) for asset_key, wiped_timestamp in wiped_timestamps_by_asset_key.items(): materialization_time = materialization_times.get(asset_key) if not materialization_time or utc_datetime_from_naive( materialization_time ) < utc_datetime_from_timestamp(wiped_timestamp): # remove rows that have not been materialized since being wiped row_by_asset_key.pop(asset_key) has_more = limit and len(rows) == limit new_cursor = rows[-1][0] if rows else None return row_by_asset_key.values(), has_more, new_cursor
def all_asset_keys(self): rows = self._fetch_asset_rows() asset_keys = [AssetKey.from_db_string(row[0]) for row in sorted(rows, key=lambda x: x[0])] return [asset_key for asset_key in asset_keys if asset_key]