<commit_before>#! /usr/bin/env python
# -*- coding: utf-8 -*-

import time
import re
from pprint import pprint
# extra libs
from neo4jrestclient.client import GraphDatabase, Node
from neo4jrestclient.query import QuerySequence
from requests.exceptions import ConnectionError
from qnib_graphdb import GraphDB
from ClusterShell.NodeSet import NodeSet

class N4jDB(GraphDB):
    def __init__(self, cfg):
        super(N4jDB, self).__init__(cfg)

    def con_gdb(self):
        """ connect to neo4j
        """
        try:
            url = "http://%(--neo4j-host)s:7474" % self._cfg
        except:
            url = "http://%(--host)s:7474" % self._cfg
        try:
            self._gdb = GraphDatabase(url)
        except ConnectionError:
            time.sleep(3)
            self.con_gdb()

    def create_node(self, node):
        """
        :param node: dict with node information

        """
        # FIXME: This is not quite right for now, since the lid is provided by the node, even though it
        # should be part of the port information, since a node could have multiple ports, hence
        # multiple lids.
        query = "MERGE (n:%(Type)s {guid:'%(Guid)s', name:'%(Name)s', lids:['%(Lid)s']}) " % node
        port_cnt = int(node['Guid'][2:], 16)
        for pnr in range(1, node['Ports'] + 1):
            port_cnt += 1
            if node['Type'] == 'SW':
                pguid = node['Guid']
            else:
                pguid = "0x%x" % port_cnt
            query += " MERGE (n)<-[:PART_OF]-(p%s:PORT {guid:'%s', port_int:'%s', lid:'%s'})" % (pnr, pguid, pnr, node['Lid'])
        self._gdb.query(query)
        msg = "Made sure node with guid '%(Guid)s' is in graph" % node
        self._cfg._logger.debug(msg)

    def merge_link(self, link, link_type):
        if link_type in ("ASICLinkList", "ASICList"):
            # Since we only have to catch SW<>SW connection we assume:
            # - looking for a SW-PORT which has the same GUID for each switch
            # - port_int determines which port to choose
            self.merge_hca_link(link)
        elif link_type in ("HCALinkList", "HCAList"):
            # a link originating from a HCA can only link to a switch...
            self.merge_sw_link(link)

    def merge_hca_link(self, link):
        """
        :param link: dict with link information
        :return:
        """
        query = "MATCH (src)<-[:PART_OF]-(sp:PORT {guid:'%(src_guid)s', port_int:'%(src_port)s'})," % link
        query += " (dest)<-[:PART_OF]-(dp:PORT {guid:'%(dst_guid)s', port_int:'%(dst_port)s'})" % link
        query += " MERGE (sp)<-[:CON]-(dp) MERGE (sp)-[:CON]->(dp)"
        query += " MERGE (src)<-[:LINKED]-(dest) MERGE (src)-[:LINKED]->(dest)"
        self._cfg._logger.debug(query)
        self._gdb.query(query, params=link)

    def merge_sw_link(self, link):
        query = "MATCH (src)<-[:PART_OF]-(sp:PORT {guid:'%(src_guid)s', port_int:'%(src_port)s'})," % link
        query += " (dest)<-[:PART_OF]-(dp:PORT {guid:'%(dst_guid)s', port_int:'%(dst_port)s'})" % link
        query += " MERGE (sp)<-[:CON]-(dp) MERGE (sp)-[:CON]->(dp)"
        query += " MERGE (src)<-[:LINKED]-(dest) MERGE (src)-[:LINKED]->(dest)"
        self._cfg._logger.debug(query)
        self._gdb.query(query, params=link)

    def merge_route(self, route):
        if route['dst_type'] == "HCA":
            # Since the destination is a HCA, we assume we are originating from an SW
            # -> Hence, ports have the same GUID, the port_int distinguishes
            query = "MATCH (dp:PORT {guid:'%(dst_guid)s'})-[:PART_OF]->(dn:%(dst_type)s)," % route
            query += " (sp:PORT {guid:'%(src_guid)s', port_int:'%(out_port)s'})" % route
            query += " MERGE (sp)-[:ROUTE_TO]->(dn)"
        else:
            # If the destination is a SW, only a SW source is of interest, since the port of an HCA is
            # only connected to one switch its the gateway to every destination
            query = "MATCH (dn:%(dst_type)s {guid:'%(dst_guid)s'})," % route
            query += " (sp:PORT {guid:'%(src_guid)s', port_int:'%(out_port)s'})" % route
            query += " MERGE (sp)-[:ROUTE_TO]->(dn)"
        self._cfg._logger.debug(query)
        self._gdb.query(query, params=route)

    def merge_slurm_partition(self, partition, nodeset):
        """
        :param partition: Name of partition
        :param nodeset: clustershell format of nodenames (e.g. compute[0-6])
        """
        query = "MATCH (n:HCA) WHERE n.name =~ '(%s).*'" % "|".join(NodeSet(nodeset))
        query += " MERGE (p:SlurmPartition {name:'%s'})" % partition
        query += " MERGE (n)-[:MEMBER]->(p)"
        self._gdb.query(query)

    def merge_slurm_job(self, jobinfo):
        """ Puts job information in cluster
        """
        query = " MERGE (j:SlurmJob {id: '%(JobId)s', name:'%(Name)s', batch: '%(BatchFlag)s', command: '%(Command)s'})" % jobinfo
        query += "ON CREATE SET j.created=timestamp(), j.state='%(JobState)s'" % jobinfo
        query += "ON MATCH SET j.seen=timestamp(), j.started='%(StartTime)s', j.state='%(JobState)s'" % jobinfo
        self._gdb.query(query)
        query = "MATCH (n:HCA) WHERE n.name =~ '(%s).*'" % "|".join(NodeSet(jobinfo['NodeList']))
        query += " MATCH (b:HCA) WHERE b.name =~ '%(BatchHost)s.*'" % jobinfo
        query += " MATCH (j:SlurmJob {id: '%(JobId)s'})" % jobinfo<commit_msg>add seen to just created jobs<commit_after>#! /usr/bin/env python
# -*- coding: utf-8 -*-

import time
import re
from pprint import pprint
# extra libs
from neo4jrestclient.client import GraphDatabase, Node
from neo4jrestclient.query import QuerySequence
from requests.exceptions import ConnectionError
from qnib_graphdb import GraphDB
from ClusterShell.NodeSet import NodeSet

class N4jDB(GraphDB):
    def __init__(self, cfg):
        super(N4jDB, self).__init__(cfg)

    def con_gdb(self):
        """ connect to neo4j
        """
        try:
            url = "http://%(--neo4j-host)s:7474" % self._cfg
        except:
            url = "http://%(--host)s:7474" % self._cfg
        try:
            self._gdb = GraphDatabase(url)
        except ConnectionError:
            time.sleep(3)
            self.con_gdb()

    def create_node(self, node):
        """
        :param node: dict with node information

        """
        # FIXME: This is not quite right for now, since the lid is provided by the node, even though it
        # should be part of the port information, since a node could have multiple ports, hence
        # multiple lids.
        query = "MERGE (n:%(Type)s {guid:'%(Guid)s', name:'%(Name)s', lids:['%(Lid)s']}) " % node
        port_cnt = int(node['Guid'][2:], 16)
        for pnr in range(1, node['Ports'] + 1):
            port_cnt += 1
            if node['Type'] == 'SW':
                pguid = node['Guid']
            else:
                pguid = "0x%x" % port_cnt
            query += " MERGE (n)<-[:PART_OF]-(p%s:PORT {guid:'%s', port_int:'%s', lid:'%s'})" % (pnr, pguid, pnr, node['Lid'])
        self._gdb.query(query)
        msg = "Made sure node with guid '%(Guid)s' is in graph" % node
        self._cfg._logger.debug(msg)

    def merge_link(self, link, link_type):
        if link_type in ("ASICLinkList", "ASICList"):
            # Since we only have to catch SW<>SW connection we assume:
            # - looking for a SW-PORT which has the same GUID for each switch
            # - port_int determines which port to choose
            self.merge_hca_link(link)
        elif link_type in ("HCALinkList", "HCAList"):
            # a link originating from a HCA can only link to a switch...
            self.merge_sw_link(link)

    def merge_hca_link(self, link):
        """
        :param link: dict with link information
        :return:
        """
        query = "MATCH (src)<-[:PART_OF]-(sp:PORT {guid:'%(src_guid)s', port_int:'%(src_port)s'})," % link
        query += " (dest)<-[:PART_OF]-(dp:PORT {guid:'%(dst_guid)s', port_int:'%(dst_port)s'})" % link
        query += " MERGE (sp)<-[:CON]-(dp) MERGE (sp)-[:CON]->(dp)"
        query += " MERGE (src)<-[:LINKED]-(dest) MERGE (src)-[:LINKED]->(dest)"
        self._cfg._logger.debug(query)
        self._gdb.query(query, params=link)

    def merge_sw_link(self, link):
        query = "MATCH (src)<-[:PART_OF]-(sp:PORT {guid:'%(src_guid)s', port_int:'%(src_port)s'})," % link
        query += " (dest)<-[:PART_OF]-(dp:PORT {guid:'%(dst_guid)s', port_int:'%(dst_port)s'})" % link
        query += " MERGE (sp)<-[:CON]-(dp) MERGE (sp)-[:CON]->(dp)"
        query += " MERGE (src)<-[:LINKED]-(dest) MERGE (src)-[:LINKED]->(dest)"
        self._cfg._logger.debug(query)
        self._gdb.query(query, params=link)

    def merge_route(self, route):
        if route['dst_type'] == "HCA":
            # Since the destination is a HCA, we assume we are originating from an SW
            # -> Hence, ports have the same GUID, the port_int distinguishes
            query = "MATCH (dp:PORT {guid:'%(dst_guid)s'})-[:PART_OF]->(dn:%(dst_type)s)," % route
            query += " (sp:PORT {guid:'%(src_guid)s', port_int:'%(out_port)s'})" % route
            query += " MERGE (sp)-[:ROUTE_TO]->(dn)"
        else:
            # If the destination is a SW, only a SW source is of interest, since the port of an HCA is
            # only connected to one switch its the gateway to every destination
            query = "MATCH (dn:%(dst_type)s {guid:'%(dst_guid)s'})," % route
            query += " (sp:PORT {guid:'%(src_guid)s', port_int:'%(out_port)s'})" % route
            query += " MERGE (sp)-[:ROUTE_TO]->(dn)"
        self._cfg._logger.debug(query)
        self._gdb.query(query, params=route)

    def merge_slurm_partition(self, partition, nodeset):
        """
        :param partition: Name of partition
        :param nodeset: clustershell format of nodenames (e.g. compute[0-6])
        """
        query = "MATCH (n:HCA) WHERE n.name =~ '(%s).*'" % "|".join(NodeSet(nodeset))
        query += " MERGE (p:SlurmPartition {name:'%s'})" % partition
        query += " MERGE (n)-[:MEMBER]->(p)"
        self._gdb.query(query)

    def merge_slurm_job(self, jobinfo):
        """ Puts job information in cluster
        """
        query = " MERGE (j:SlurmJob {id: '%(JobId)s', name:'%(Name)s', batch: '%(BatchFlag)s', command: '%(Command)s'})" % jobinfo
        query += "ON CREATE SET j.created=timestamp(), j.seen=timestamp(), j.state='%(JobState)s'" % jobinfo
        query += "ON MATCH SET j.seen=timestamp(), j.started='%(StartTime)s', j.state='%(JobState)s'" % jobinfo
        self._gdb.query(query)
        query = "MATCH (n:HCA) WHERE n.name =~ '(%s).*'" % "|".join(NodeSet(jobinfo['NodeList']))
        query += " MATCH (b:HCA) WHERE b.name =~ '%(BatchHost)s.*'" % jobinfo
        query += " MATCH (j:SlurmJob {id: '%(JobId)s'})" % jobinfo
{'fraction_non_alphanumeric': 0.17264532744665195, 'fraction_numerical': 0.00459896983075791, 'mean_word_length': 2.5860817941952505, 'pattern_counts': {'":': 2, '<': 23, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '937556', 'n_tokens_mistral': 3764, 'n_tokens_neox': 3617, 'n_words': 911}