def _calculate_outputs(
topo_traversal: GraphTraversalSource,
source_vertex: Vertex,
arrival_rates: ARRIVAL_RATES,
output_rates: DefaultDict[int, Dict[str, float]],
coefficients: pd.Series,
) -> DefaultDict[int, Dict[str, float]]:
source_task: int = (
topo_traversal.V(source_vertex).properties("task_id").value().next())
in_streams: List[Dict[str, str]] = \
(topo_traversal.V(source_vertex).inE("logically_connected")
.project("stream_name", "source_component")
.by(properties("stream").value())
.by(outV().properties("component").value())
.dedup()
.toList())
out_streams: List[str] = \
(topo_traversal.V(source_vertex)
.outE("logically_connected").values("stream")
.dedup().toList())
for out_stream in out_streams:
output_rate: float = 0.0
for in_stream in in_streams:
in_stream_name: str = in_stream["stream_name"]
source_component: str = in_stream["source_component"]
stream_arrivals: float = \
arrival_rates[source_task][(in_stream_name,
source_component)]
try:
coefficent: float = float(coefficients.loc[source_task,
out_stream,
in_stream_name,
source_component])
except KeyError:
LOG.debug(
"No coefficient available for source task %d, "
"out stream %s, in stream %s from component %s",
source_task, out_stream, in_stream_name, source_component)
else:
output_rate += (stream_arrivals * coefficent)
# It is possible that some of the IO coefficients may be negative,
# implying that the more you receive on an input stream the less you
# output to a given output stream. If we anticipate a large arrival on
# this negative input stream and low on other positive streams then it
# is possible that the predicted output rate could be negative (which
# is obviously meaningless).
if output_rate < 0.0:
output_rate = 0.0
output_rates[source_task][out_stream] = output_rate
return output_rates
def _expand_properties(g: GraphTraversalSource, column: str,
value: Any) -> GraphTraversalSource:
# If this is a list then expand it out into multiple property calls
if isinstance(value, list) and len(value) > 0:
for item in value:
g = g.property(Cardinality.set_, column, item)
elif pd.notna(value):
g = g.property(Cardinality.set_, column, value)
return g
def _build_gremlin_properties(g: GraphTraversalSource,
row: Any) -> GraphTraversalSource:
for (column, value) in row.items():
if column not in ["~id", "~label", "~to", "~from"]:
if isinstance(value, list) and len(value) > 0:
for item in value:
g = g.property(Cardinality.set_, column, item)
elif not pd.isna(value) and not pd.isnull(value):
g = g.property(column, value)
return g
def _user_search_query(graph: GraphTraversalSource,
tag_filter: str) -> List[Dict]:
traversal = graph.V().hasLabel(User.USER_NODE_LABEL)
traversal = traversal.has(User.USER_NODE_FULL_NAME)
if tag_filter:
traversal = traversal.where('published_tag', tag_filter)
traversal = traversal.project('email', 'first_name', 'last_name',
'full_name', 'github_username', 'team_name',
'employee_type', 'manager_email', 'slack_id',
'is_active', 'role_name', 'total_read',
'total_own', 'total_follow')
traversal = traversal.by('email') # email
traversal = traversal.by('first_name') # first_name
traversal = traversal.by('last_name') # last_name
traversal = traversal.by('full_name') # full_name
traversal = traversal.by('github_username') # github_username
traversal = traversal.by('team_name') # team_name
traversal = traversal.by('employee_type') # employee_type
traversal = traversal.by(
__.coalesce(
__.out(User.USER_MANAGER_RELATION_TYPE).values('email'),
__.constant(''))) # manager_email
traversal = traversal.by('slack_id') # slack_id
traversal = traversal.by('is_active') # is_active
traversal = traversal.by('role_name') # role_name
traversal = traversal.by(
__.coalesce(
__.outE(READ_RELATION_TYPE).values('read_count'),
__.constant(0)).sum()) # total_read
traversal = traversal.by(
__.outE(OWNER_OF_OBJECT_RELATION_TYPE).fold().count()) # total_own
traversal = traversal.by(
__.outE('FOLLOWED_BY').fold().count()) # total_follow
traversal = traversal.order().by(__.select('email'), Order.asc)
return traversal.toList()
def get_comp_links_by_grouping(graph_traversal: GraphTraversalSource,
grouping: str) -> List[Dict[str, str]]:
""" Gets a list of component connection dictionaries. These describe all
source->stream->destination connections with the specified grouping value
in the topology available via the supplied graph traversal source.
Arguments:
graph_traversal (GraphTraversalSource): A GraphTraversalSource instance
linked to the topology subgraph
whose connections are to be
queried.
grouping (str): The stream grouping of the connections to be returned.
Returns:
A list of dictionaries each containing "source", "stream" and
"destination" keys of the component and stream name respectively.
"""
component_connections: List[Dict[str, str]] = \
(graph_traversal.V().hasLabel(P.within("bolt", "spout")).as_("source")
.outE("logically_connected").has("grouping", grouping).as_("stream")
.inV().as_("destination").select("source", "stream", "destination")
.by("component").by("stream").by("component").dedup().toList())
return component_connections
def _table_search_query(graph: GraphTraversalSource,
tag_filter: str) -> List[Dict]:
traversal = graph.V().hasLabel(TableMetadata.TABLE_NODE_LABEL)
if tag_filter:
traversal = traversal.has('published_tag', tag_filter)
traversal = traversal.project('database', 'cluster', 'schema',
'schema_description', 'name', 'key',
'description', 'last_updated_timestamp',
'column_names', 'column_descriptions',
'total_usage', 'unique_usage', 'tags',
'badges', 'programmatic_descriptions')
traversal = traversal.by(
__.out(TableMetadata.TABLE_SCHEMA_RELATION_TYPE).out(
SCHEMA_REVERSE_RELATION_TYPE).out(
CLUSTER_REVERSE_RELATION_TYPE).values('name')) # database
traversal = traversal.by(
__.out(TableMetadata.TABLE_SCHEMA_RELATION_TYPE).out(
SCHEMA_REVERSE_RELATION_TYPE).values('name')) # cluster
traversal = traversal.by(
__.out(
TableMetadata.TABLE_SCHEMA_RELATION_TYPE).values('name')) # schema
traversal = traversal.by(
__.coalesce(
__.out(TableMetadata.TABLE_SCHEMA_RELATION_TYPE).out(
DescriptionMetadata.DESCRIPTION_RELATION_TYPE).values(
'description'), __.constant(''))) # schema_description
traversal = traversal.by('name') # name
traversal = traversal.by(T.id) # key
traversal = traversal.by(
__.coalesce(
__.out(DescriptionMetadata.DESCRIPTION_RELATION_TYPE).values(
'description'), __.constant(''))) # description
traversal = traversal.by(
__.coalesce(
__.out(LASTUPDATED_RELATION_TYPE).values(TIMESTAMP_PROPERTY),
__.constant(''))) # last_updated_timestamp
traversal = traversal.by(
__.out(TableMetadata.TABLE_COL_RELATION_TYPE).values(
'name').fold()) # column_names
traversal = traversal.by(
__.out(TableMetadata.TABLE_COL_RELATION_TYPE).out(
DescriptionMetadata.DESCRIPTION_RELATION_TYPE).values(
'description').fold()) # column_descriptions
traversal = traversal.by(
__.coalesce(
__.outE(READ_REVERSE_RELATION_TYPE).values('read_count'),
__.constant(0)).sum()) # total_usage
traversal = traversal.by(
__.outE(READ_REVERSE_RELATION_TYPE).count()) # unique_usage
traversal = traversal.by(
__.inE(TableMetadata.TAG_TABLE_RELATION_TYPE).outV().values(
METADATA_KEY_PROPERTY_NAME).fold()) # tags
traversal = traversal.by(
__.out('HAS_BADGE').values('keys').dedup().fold()) # badges
traversal = traversal.by(
__.out(DescriptionMetadata.PROGRAMMATIC_DESCRIPTION_NODE_LABEL).values(
'description').fold()) # programmatic_descriptions
traversal = traversal.order().by(__.select('name'), Order.asc)
return traversal.toList()
def _build_gremlin_vertices(g: GraphTraversalSource,
row: Any) -> GraphTraversalSource:
g = g.V(str(row["~id"])).fold().coalesce(
__.unfold(),
__.addV(row["~label"]).property(T.id, str(row["~id"])))
g = _build_gremlin_properties(g, row)
return g
def _build_gremlin_insert_vertices(
g: GraphTraversalSource,
row: Any,
use_header_cardinality: bool = False) -> GraphTraversalSource:
g = g.V(str(row["~id"])).fold().coalesce(
__.unfold(),
__.addV(row["~label"]).property(T.id, str(row["~id"])))
g = _set_properties(g, use_header_cardinality, row)
return g
def expire_connections_for_other(cls, *, _g: GraphTraversalSource,
vertex_type: VertexType,
keys: FrozenSet[str],
existing: EXISTING) -> None:
# V().has(label, 'key', P.without(keys)) is more intuitive but doesn't scale, so instead just find all those
g = _g.V().hasLabel(vertex_type.label).where(__.bothE())
g = g.values(WellKnownProperties.Key.value.name)
all_to_expire_keys = set(g.toList()).difference(keys)
# TODO: when any vertex ids that need something besides key
all_to_expire = set(
vertex_type.id(key=key) for key in all_to_expire_keys)
for to_expire in chunk(all_to_expire, 1000):
g = _g.V(tuple(to_expire)).bothE()
g = g.local(
__.union(__.outV().id(), __.valueMap(True),
__.inV().id()).fold())
cls._into_existing(g.toList(), existing)
def _set_properties(g: GraphTraversalSource, use_header_cardinality: bool,
row: Any) -> GraphTraversalSource:
for (column, value) in row.items():
if column not in ["~id", "~label", "~to", "~from"]:
# If the column header is specifying the cardinality then use it
if use_header_cardinality:
if column.lower().find("(single)") > 0 and pd.notna(value):
g = g.property(Cardinality.single,
_get_column_name(column), value)
else:
g = _expand_properties(g, _get_column_name(column), value)
else:
# If not using header cardinality then use the default of set
g = _expand_properties(g, column, value)
return g
def _build_gremlin_edges(g: GraphTraversalSource,
row: pd.Series) -> GraphTraversalSource:
g = (g.V(str(row["~from"])).fold().coalesce(
__.unfold(),
_build_gremlin_vertices(__, {
"~id": row["~from"],
"~label": "Vertex"
})).addE(row["~label"]).to(
__.V(str(row["~to"])).fold().coalesce(
__.unfold(),
_build_gremlin_vertices(__, {
"~id": row["~to"],
"~label": "Vertex"
}))))
g = _build_gremlin_properties(g, row)
return g
def _column_entities(cls, *, _g: GraphTraversalSource,
tables_ids: Iterable[str],
existing: EXISTING) -> None:
# fetch database -> cluster -> schema -> table links
g = _g.V(tuple(tables_ids))
g = g.outE(EdgeTypes.Column.value.label)
g = g.inV().hasLabel(VertexTypes.Column.value.label).as_('columns')
# fetch column -> links (no Stat)
for t in [EdgeTypes.Description]:
g = g.coalesce(__.select('columns').outE(
t.value.label).fold()).as_(t.name)
g = g.select(EdgeTypes.Description.name).unfold()
g = g.local(
__.union(__.outV().id(), __.valueMap(True),
__.inV().id()).fold())
cls._into_existing(g.toList(), existing)
def _build_gremlin_insert_edges(
g: GraphTraversalSource, row: pd.Series,
use_header_cardinality: bool) -> GraphTraversalSource:
g = (g.V(str(row["~from"])).fold().coalesce(
__.unfold(),
_build_gremlin_insert_vertices(__, {
"~id": row["~from"],
"~label": "Vertex"
})).addE(row["~label"]).property(T.id, str(row["~id"])).to(
__.V(str(row["~to"])).fold().coalesce(
__.unfold(),
_build_gremlin_insert_vertices(__, {
"~id": row["~to"],
"~label": "Vertex"
}))))
g = _set_properties(g, use_header_cardinality, row)
return g
def traversal(self, *, graph_traversal=None, remote_strategy=None):
"""
Get a traversal source from the Graph
:param gremlin_python.process.GraphTraversal graph_traversal:
Custom graph traversal class
:param gremlin_python.driver.remote_connection.RemoteStrategy remote_strategy:
Custom remote strategy class
:returns:
:py:class:`gremlin_python.process.graph_traversal.GraphTraversalSource`
"""
if graph_traversal is None:
graph_traversal = AsyncGraphTraversal
if remote_strategy is None:
remote_strategy = AsyncRemoteStrategy
return GraphTraversalSource(
self,
TraversalStrategies.global_cache[self.__class__],
remote_strategy=remote_strategy,
graph_traversal=graph_traversal)
def get_levels(topo_traversal: GraphTraversalSource) -> List[List[Vertex]]:
""" Gets the levels of the logical graph. The traversal starts with the
source spouts and performs a breadth first search through the logically
connected vertices.
Arguments:
topo_traversal (GraphTraversalSource): A traversal source instance
mapped to the topology subgraph
whose levels are to be
calculated.
Returns:
A list where each entry is a list of Vertex instances representing a
level within the logical graph. The first level will be the spout
instances.
"""
# Only load the static enums we need so we don't pollute the globals dict
keys = statics.staticEnums["keys"]
values = statics.staticEnums["values"]
local_scope = statics.staticEnums["local"]
# Repeatedly traverse the tree defined by the logical connections, grouping
# each group (or set because we us de-duplicate) of vertices by their depth
# in the tree. This depth is based the current number of times the repeat
# step has run (loops). So you end up with a map of integer depth to list
# of vertices which is emitted by the cap step. After this we just put the
# Hash Map in key order (ascending) and then take only the values (the
# lists of vertices) and unfold them into a list.
# The first group by(-1) statement is so that the spout vertices are
# included at the top of the list
levels: List[List[Vertex]] = (
topo_traversal.V().hasLabel("spout").group("m").by(
constant(-1)).repeat(
out("logically_connected").dedup().group("m").by(
loops())).until(not_(outE("logically_connected"))).
cap("m").order(local_scope).by(keys).select(values).unfold().toList())
return levels
def _build_gremlin_update(
g: GraphTraversalSource, row: Any,
use_header_cardinality: bool) -> GraphTraversalSource:
g = g.V(str(row["~id"]))
g = _set_properties(g, use_header_cardinality, row)
return g
def _calculate_arrivals(topo_traversal: GraphTraversalSource,
source_vertex: Vertex, arrival_rates: ARRIVAL_RATES,
output_rates: DefaultDict[int, Dict[str, float]],
i2i_rps: pd.DataFrame) -> ARRIVAL_RATES:
# Get all downstream edges and vertices for this source vertex
out_edges: List[Dict[str, Union[str, int, float]]] = \
(topo_traversal.V(source_vertex).outE("logically_connected")
.project("source_task", "source_component", "stream_name",
"destination_task", "destination_component")
.by(outV().properties("task_id").value())
.by(outV().properties("component").value())
.by(properties("stream").value())
.by(inV().properties("task_id").value())
.by(inV().properties("component").value())
.toList())
if not out_edges:
return arrival_rates
source_task: int = cast(int, out_edges[0]["source_task"])
source_component: str = cast(str, out_edges[0]["source_component"])
LOG.debug("Processing output from source instance %s_%d", source_component,
source_task)
for out_edge in out_edges:
stream: str = cast(str, out_edge["stream_name"])
try:
stream_output: float = cast(float,
output_rates[source_task][stream])
except KeyError:
LOG.debug(
"No output rate information for source task %d on "
"stream %s. Skipping the outgoing edge", source_task, stream)
continue
destination_task: int = cast(int, out_edge["destination_task"])
try:
r_prob: float = float(
i2i_rps.loc(axis=0)[source_task, destination_task, stream])
except KeyError:
LOG.debug(
"Unable to find routing probability for connection from "
"task %d to %d on stream %s", source_task, destination_task,
stream)
edge_output: float = 0.0
else:
edge_output = (stream_output * r_prob)
LOG.debug(
"Output from %s-%d to %s-%d on stream %s is "
"calculated as %f * %f = %f", source_component, source_task,
out_edge["destination_component"], destination_task, stream,
stream_output, r_prob, edge_output)
arrival_rates[destination_task][(stream,
source_component)] += edge_output
return arrival_rates
def delete_graph_for_shard_only(g: GraphTraversalSource) -> None:
shard = get_shard()
assert shard, f'expected shard to exist! Surely you are only using this in development or test?'
# TODO: do something better than not using WellKnownProperties.TestShard here (since that makes a circular
# dependency)
g.V().has('shard', shard).drop().iterate()
def _dashboard_search_query(graph: GraphTraversalSource,
tag_filter: str) -> List[Dict]:
traversal = graph.V().hasLabel(DashboardMetadata.DASHBOARD_NODE_LABEL)
traversal = traversal.has('name')
if tag_filter:
traversal = traversal.where('published_tag', tag_filter)
traversal = traversal.project('group_name', 'name', 'cluster',
'description', 'group_description',
'group_url', 'url', 'uri',
'last_successful_run_timestamp',
'query_names', 'chart_names', 'total_usage',
'tags', 'badges')
traversal = traversal.by(
__.out(
DashboardMetadata.DASHBOARD_DASHBOARD_GROUP_RELATION_TYPE).values(
'name')) # group_name
traversal = traversal.by('name') # name
traversal = traversal.by(
__.out(DashboardMetadata.DASHBOARD_DASHBOARD_GROUP_RELATION_TYPE).out(
DashboardMetadata.DASHBOARD_GROUP_CLUSTER_RELATION_TYPE).values(
'name')) # cluster
traversal = traversal.by(
__.coalesce(
__.out(
DashboardMetadata.DASHBOARD_DESCRIPTION_RELATION_TYPE).values(
'description'), __.constant(''))) # description
traversal = traversal.by(
__.coalesce(
__.out(
DashboardMetadata.DASHBOARD_DASHBOARD_GROUP_RELATION_TYPE).out(
DashboardMetadata.DASHBOARD_DESCRIPTION_RELATION_TYPE).
values('description'), __.constant(''))) # group_description
traversal = traversal.by(
__.out(
DashboardMetadata.DASHBOARD_DASHBOARD_GROUP_RELATION_TYPE).values(
'dashboard_group_url')) # group_url
traversal = traversal.by('dashboard_url') # dashboard_url
traversal = traversal.by('key') # uri
traversal = traversal.by(
__.coalesce(
__.out('EXECUTED').has(
'key', TextP.endingWith(
'_last_successful_execution')).values('timestamp'),
__.constant(''))) # last_successful_run_timestamp
traversal = traversal.by(
__.out(DashboardQuery.DASHBOARD_QUERY_RELATION_TYPE).values(
'name').dedup().fold()) # query_names
traversal = traversal.by(
__.out(DashboardQuery.DASHBOARD_QUERY_RELATION_TYPE).out(
DashboardChart.CHART_RELATION_TYPE).values(
'name').dedup().fold()) # chart_names
traversal = traversal.by(
__.coalesce(
__.outE(READ_REVERSE_RELATION_TYPE).values(
READ_RELATION_COUNT_PROPERTY),
__.constant(0)).sum()) # total_usage
traversal = traversal.by(
__.out('TAGGED_BY').has(
'tag_type', 'default').values('keys').dedup().fold()) # tags
traversal = traversal.by(
__.out('HAS_BADGE').values('keys').dedup().fold()) # badges
traversal = traversal.order().by(__.select('name'), Order.asc)
dashboards = traversal.toList()
for dashboard in dashboards:
dashboard['product'] = dashboard['uri'].split('_')[0]
return dashboards
def delete_everything(traversal: GraphTraversalSource):
return traversal.V().drop().toList()
def traversal(self):
return GraphTraversalSource(self, TraversalStrategies.global_cache[self.__class__])
def _build_gremlin_update(g: GraphTraversalSource,
row: Any) -> GraphTraversalSource:
g = g.V(str(row["~id"]))
g = _build_gremlin_properties(g, row)
return g
def table_entities(cls, *, _g: GraphTraversalSource,
table_data: List[Table], existing: EXISTING) -> None:
all_tables_ids = list(
set([
VertexTypes.Table.value.id(
key=TableUris.get(database=t.database,
cluster=t.cluster,
schema=t.schema,
table=t.name).table) for t in table_data
]))
all_owner_ids = list(
set([
VertexTypes.User.value.id(key=key) for key in [
t.table_writer.id for t in table_data
if t.table_writer is not None
]
]))
all_application_ids = list(
set(
list(
possible_vertex_ids_for_application_key(*[
t.table_writer.id for t in table_data
if t.table_writer is not None
]))))
# chunk these since 100,000s seems to choke
for tables_ids in chunk(all_tables_ids, 1000):
LOGGER.info(f'fetching for tables: {tables_ids}')
# fetch database -> cluster -> schema -> table links
g = _g.V(tuple(tables_ids)).as_('tables')
g = g.coalesce(__.inE(
EdgeTypes.Table.value.label).dedup().fold()).as_(
EdgeTypes.Table.name)
g = g.coalesce(__.unfold().outV().hasLabel(
VertexTypes.Schema.value.label).inE(
EdgeTypes.Schema.value.label).dedup().fold()).as_(
EdgeTypes.Schema.name)
g = g.coalesce(__.unfold().outV().hasLabel(
VertexTypes.Cluster.value.label).inE(
EdgeTypes.Cluster.value.label).dedup().fold()).as_(
EdgeTypes.Cluster.name)
# fetch table <- links
for t in (EdgeTypes.BelongToTable, EdgeTypes.Generates,
EdgeTypes.Tag):
g = g.coalesce(__.select('tables').inE(
t.value.label).fold()).as_(t.name)
# fetch table -> column et al links
for t in (EdgeTypes.Column, EdgeTypes.Description,
EdgeTypes.LastUpdatedAt, EdgeTypes.Source,
EdgeTypes.Stat):
g = g.coalesce(__.select('tables').outE(
t.value.label).fold()).as_(t.name)
# TODO: add owners, watermarks, last timestamp existing, source
aliases = set([
t.name
for t in (EdgeTypes.Table, EdgeTypes.Schema, EdgeTypes.Cluster,
EdgeTypes.BelongToTable, EdgeTypes.Generates,
EdgeTypes.Tag, EdgeTypes.Column,
EdgeTypes.Description, EdgeTypes.LastUpdatedAt,
EdgeTypes.Source, EdgeTypes.Stat)
])
g = g.select(*aliases).unfold().select(MapColumn.values).unfold()
g = g.local(
__.union(__.outV().id(), __.valueMap(True),
__.inV().id()).fold())
cls._into_existing(g.toList(), existing)
cls._column_entities(_g=_g,
tables_ids=tables_ids,
existing=existing)
# fetch Application, User
for ids in chunk(list(set(all_application_ids + all_owner_ids)), 5000):
LOGGER.info(f'fetching for application/owners: {ids}')
g = _g.V(ids).valueMap(True)
cls._into_existing(g.toList(), existing)