From 5ea92399d65a49a608868598035c4e64611a9735 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleksa=20Elezovi=C4=87?= <aleksaelezovic@gmail.com>
Date: Thu, 28 Nov 2024 11:27:12 +0100
Subject: [PATCH] implement graph.query

---
 dkg/graph.py              | 16 +++++++++++-----
 dkg/types/__init__.py     |  2 +-
 dkg/types/dkg_node.py     | 15 +++++++++++++++
 dkg/utils/node_request.py | 10 ++++++++--
 4 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/dkg/graph.py b/dkg/graph.py
index 43cf234..049dd52 100644
--- a/dkg/graph.py
+++ b/dkg/graph.py
@@ -22,7 +22,7 @@
 from dkg.manager import DefaultRequestManager
 from dkg.method import Method
 from dkg.module import Module
-from dkg.types import NQuads
+from dkg.types import NQuads, GraphState, GraphLocation, UAL
 from dkg.utils.decorators import retry
 from dkg.utils.node_request import NodeRequest, validate_operation_status
 
@@ -37,14 +37,20 @@ def __init__(self, manager: DefaultRequestManager):
     def query(
         self,
         query: str,
-        repository: str,
+        graphState: GraphState = GraphState.CURRENT,
+        graphLocation: GraphLocation = GraphLocation.PUBLIC,
+        paranetUAL: UAL = "",
     ) -> NQuads:
         parsed_query = parseQuery(query)
         query_type = parsed_query[1].name.replace("Query", "").upper()
 
-        operation_id: NodeResponseDict = self._query(query, query_type, repository)[
-            "operationId"
-        ]
+        operation_id: NodeResponseDict = self._query(
+            query,
+            query_type,
+            graphState,
+            graphLocation,
+            paranetUAL,
+        )["operationId"]
         operation_result = self.get_operation_result(operation_id, "query")
 
         return operation_result["data"]
diff --git a/dkg/types/__init__.py b/dkg/types/__init__.py
index 416728b..f02c57e 100644
--- a/dkg/types/__init__.py
+++ b/dkg/types/__init__.py
@@ -13,7 +13,7 @@
     AgreementData,
     Environment,
 )
-from .dkg_node import UAL  # NOQA: F401
+from .dkg_node import UAL, GraphState, GraphLocation  # NOQA: F401
 from .encoding import BytesLike, DataHexStr, HexStr  # NOQA: F401
 from .evm import Address, ChecksumAddress, Wei  # NOQA: F401
 from .generics import TFunc, TReturn  # NOQA: F401
diff --git a/dkg/types/dkg_node.py b/dkg/types/dkg_node.py
index 02a8e62..a85968c 100644
--- a/dkg/types/dkg_node.py
+++ b/dkg/types/dkg_node.py
@@ -16,5 +16,20 @@
 # under the License.
 
 from typing import NewType
+from enum import Enum
 
 UAL = NewType("UAL", str)
+
+class GraphLocation(str, Enum):
+    PUBLIC = 'PUBLIC_KG'
+    PRIVATE = 'LOCAL_KG'
+
+    def __str__(self) -> str:
+        return self.value
+
+class GraphState(str, Enum):
+    CURRENT = 'CURRENT'
+    HISTORICAL = 'HISTORICAL'
+
+    def __str__(self) -> str:
+        return self.value
\ No newline at end of file
diff --git a/dkg/utils/node_request.py b/dkg/utils/node_request.py
index 977dc8e..6fa5f3f 100644
--- a/dkg/utils/node_request.py
+++ b/dkg/utils/node_request.py
@@ -21,7 +21,7 @@
 
 from dkg.dataclasses import BidSuggestionRange, HTTPRequestMethod
 from dkg.exceptions import OperationFailed, OperationNotFinished
-from dkg.types import AutoStrEnumUpperCase, UAL, Address, DataHexStr, NQuads
+from dkg.types import AutoStrEnumUpperCase, UAL, Address, DataHexStr, NQuads, GraphState, GraphLocation
 
 
 @dataclass
@@ -77,7 +77,13 @@ class NodeRequest:
     query = NodeCall(
         method=HTTPRequestMethod.POST,
         path="query",
-        data={"query": str, "type": str, "repository": str},
+        data={
+            "query": str,
+            "type": str,
+            "graphState": GraphState,
+            "graphLocation": GraphLocation,
+            "paranetUAL": UAL,
+        },
     )