GraphQLSender.py

import struct
from utils import pack_string
from ParseResponse import unpack_packstream, clean_graph_node

def recv_all(socket, bits):
    """
    A helper method used to process all received bytes and treat them as a single return.
    :param socket: The socket containing a returned message.
    :param bits: the number of bits to read at once.
    :return: the received bytes.
    """
    data = bytearray()
    while len(data) < bits:
        packet = socket.recv(bits-len(data))
        if not packet:
            raise ConnectionError("Socket closed by server")
        data.extend(packet)
    return bytes(data)

def run_cypher_query(s, cypher_string):
    """
    Converts the provided string into bytes and formats them appropriately for the database.
    :param s: the socket used to send the query.
    :param cypher_string: A string containing CypherQL to run against the database.
    :return: A list of dictionary objects.
    """
    query_param_map = bytes([0xA0]) ## empty map.

    run_body = bytes([0xB2, 0x10]) + pack_string(cypher_string) + query_param_map

    run_chunk_header = struct.pack('>H', len(run_body))
    run_packet = run_chunk_header + run_body + b'\x00\x00'

    pull_body = bytes([0xB0, 0x3F])

    pull_chunk_header = struct.pack('>H', len(pull_body))
    pull_packet = pull_chunk_header + pull_body + b'\x00\x00'

    s.sendall(run_packet + pull_packet)

    resp_header = recv_all(s, 2)
    chunk_size = struct.unpack('>H', resp_header)[0]
    resp_body = recv_all(s, chunk_size)
    recv_all(s, 2) ## consume end of message marker

    decoded_data, _ = unpack_packstream(resp_body)
    header_data = []
    if decoded_data['_type'] == "STRUCT_0x70":
        header_data = decoded_data['_fields'][0]['fields']
    elif decoded_data['_type'] == "STRUCT_0x7f": ##ERROR RECEIVED.
        raise Exception(f"[DATABASE ERROR]: {decoded_data['_fields']}")

    record_count = 0
    results = []
    while True:
        resp_header = recv_all(s, 2)
        chunk_size = struct.unpack('>H', resp_header)[0]

        if chunk_size == 0:
            break

        resp_body = recv_all(s, chunk_size)
        recv_all(s, 2)

        decoded_data, _ = unpack_packstream(resp_body)
        msg_type = decoded_data['_type']

        if msg_type == "STRUCT_0x71":
            record_count += 1
            raw_rows = decoded_data['_fields']
            cleaned_result = clean_graph_node(raw_rows, header_data)
            results.append(cleaned_result)
        elif msg_type == "STRUCT_0x7f":
            raise Exception(f"[DATABASE ERROR]: {decoded_data['_fields']}")
        elif msg_type == "STRUCT_0x70":
            break
    return results