Inject Arrow credentials to UploadGraph endpoint

Co-authored-by: Brian Shi <brian.shi@neotechnology.com>
Co-authored-by: Olga Razvenskaia <olga.razvenskaia@neo4j.com>

Author: 3 people

examples/python-runtime.ipynb

@@ -3,9 +3,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "collapsed": true
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "DBID = \"beefbeef\"\n",
@@ -14,68 +12,46 @@
     "\n",
     "from graphdatascience import GraphDataScience\n",
     "\n",
-    "gds = GraphDataScience(\n",
-    "    f\"neo4j+s://{DBID}-{ENVIRONMENT}.databases.neo4j-dev.io/\", auth=(\"neo4j\", PASSWORD)\n",
-    ")\n",
+    "gds = GraphDataScience(f\"neo4j+s://{DBID}-{ENVIRONMENT}.databases.neo4j-dev.io/\", auth=(\"neo4j\", PASSWORD))\n",
     "gds.set_database(\"neo4j\")\n",
     "\n",
-    "gds.gnn.nodeClassification.train(\"cora\", \"model\", [\"features\"], \"subject\", node_labels=[\"Paper\"])\n"
+    "gds.gnn.nodeClassification.train(\"cora\", \"model\", [\"features\"], \"subject\", node_labels=[\"Paper\"])"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "metadata": {},
    "outputs": [],
    "source": [
     "try:\n",
     "    gds.graph.load_cora()\n",
     "except:\n",
-    "    pass\n"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
+    "    pass"
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "metadata": {},
    "outputs": [],
    "source": [
-    "gds.gnn.nodeClassification.train(\"cora\", \"model\", [\"features\"], \"subject\", node_labels=[\"Paper\"])\n"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
+    "gds.gnn.nodeClassification.train(\"cora\", \"model\", [\"features\"], \"subject\", node_labels=[\"Paper\"])"
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "metadata": {},
    "outputs": [],
    "source": [
     "gds.gnn.nodeClassification.predict(\"cora\", \"model\", [\"features\"], \"subject\", node_labels=[\"Paper\"])"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
+   ]
   }
  ],
  "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  },
   "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 2
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython2",
-   "version": "2.7.6"
+   "name": "python"
   }
  },
  "nbformat": 4,

graphdatascience/endpoints.py

@@ -33,7 +33,9 @@
 """
 
 
-class DirectEndpoints(DirectSystemEndpoints, DirectUtilEndpoints, GraphEndpoints, PipelineEndpoints, ModelEndpoints, GnnEndpoints):
+class DirectEndpoints(
+    DirectSystemEndpoints, DirectUtilEndpoints, GraphEndpoints, PipelineEndpoints, ModelEndpoints, GnnEndpoints
+):
     def __init__(self, query_runner: QueryRunner, namespace: str, server_version: ServerVersion):
         super().__init__(query_runner, namespace, server_version)
 

graphdatascience/gnn/gnn_endpoints.py

@@ -1,17 +1,18 @@
-from .gnn_nc_runner import GNNNodeClassificationRunner
 from ..caller_base import CallerBase
 from ..error.illegal_attr_checker import IllegalAttrChecker
 from ..error.uncallable_namespace import UncallableNamespace
+from .gnn_nc_runner import GNNNodeClassificationRunner
+
 
 class GNNRunner(UncallableNamespace, IllegalAttrChecker):
     @property
     def nodeClassification(self) -> GNNNodeClassificationRunner:
-        return GNNNodeClassificationRunner(self._query_runner, f"{self._namespace}.nodeClassification", self._server_version)
+        return GNNNodeClassificationRunner(
+            self._query_runner, f"{self._namespace}.nodeClassification", self._server_version
+        )
+
 
 class GnnEndpoints(CallerBase):
     @property
     def gnn(self) -> GNNRunner:
         return GNNRunner(self._query_runner, f"{self._namespace}.gnn", self._server_version)
-
-
-

graphdatascience/gnn/gnn_nc_runner.py

@@ -13,13 +13,24 @@ def train(self, graph_name: str, model_name: str, feature_properties: List[str],
             "targetProperty": target_property,
             "job_type": "train",
         }
+
         node_properties = feature_properties + [target_property]
         if target_node_label:
             configMap["targetNodeLabel"] = target_node_label
         mlTrainingConfig = json.dumps(configMap)
         # TODO query available node labels
         node_labels = ["Paper"] if not node_labels else node_labels
-        self._query_runner.run_query(f"CALL gds.upload.graph('{graph_name}', {{mlTrainingConfig: '{mlTrainingConfig}', modelName: '{model_name}', nodeLabels: {node_labels}, nodeProperties: {node_properties}}})")
+
+        # token and uri will be injected by arrow_query_runner
+        self._query_runner.run_query(
+            f"CALL gds.upload.graph($graph_name, $config)", 
+            params={"graph_name": graph_name, "config": {
+                     "mlTrainingConfig": mlTrainingConfig, 
+                     "modelName": model_name, 
+                     "nodeLabels": node_labels,
+                    "nodeProperties": node_properties
+                    }}
+                    )
 
 
     def predict(self, graph_name: str, model_name: str, feature_properties: List[str], target_node_label: str = None, node_labels: List[str] = None) -> "Series[Any]":

graphdatascience/query_runner/arrow_query_runner.py

@@ -29,6 +29,9 @@ def __init__(
     ):
         self._fallback_query_runner = fallback_query_runner
         self._server_version = server_version
+        # FIXME handle version were tls cert is given
+        self._auth = auth
+        self._uri = uri
 
         host, port_string = uri.split(":")
 
@@ -39,8 +42,9 @@ def __init__(
         )
 
         client_options: Dict[str, Any] = {"disable_server_verification": disable_server_verification}
+        self._auth_factory = AuthFactory(auth)
         if auth:
-            client_options["middleware"] = [AuthFactory(auth)]
+            client_options["middleware"] = [self._auth_factory]
         if tls_root_certs:
             client_options["tls_root_certs"] = tls_root_certs
 
@@ -129,6 +133,11 @@ def run_query(
                 endpoint = "gds.beta.graph.relationships.stream"
 
             return self._run_arrow_property_get(graph_name, endpoint, {"relationship_types": relationship_types})
+        elif "gds.upload.graph" in query:
+            # inject parameters
+            params["config"]["token"] = self._get_or_request_token()
+            params["config"]["arrowEndpoint"] = self._uri
+            print(params)
 
         return self._fallback_query_runner.run_query(query, params, database, custom_error)
 
@@ -183,11 +192,19 @@ def create_graph_constructor(
         return ArrowGraphConstructor(
             database, graph_name, self._flight_client, concurrency, undirected_relationship_types
         )
+    
+    def _get_or_request_token(self) -> str:
+        print("get or request token")
+        self._flight_client.authenticate_basic_token(self._auth[0], self._auth[1])
+        return self._auth_factory.token()
 
 
 class AuthFactory(ClientMiddlewareFactory):  # type: ignore
     def __init__(self, auth: Tuple[str, str], *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
+        print("init auth factory")
+
+
         self._auth = auth
         self._token: Optional[str] = None
         self._token_timestamp = 0
@@ -196,6 +213,7 @@ def start_call(self, info: Any) -> "AuthMiddleware":
         return AuthMiddleware(self)
 
     def token(self) -> Optional[str]:
+        print(f"current token {self._token} at {self._token_timestamp}")
         # check whether the token is older than 10 minutes. If so, reset it.
         if self._token and int(time.time()) - self._token_timestamp > 600:
             self._token = None
@@ -206,6 +224,8 @@ def set_token(self, token: str) -> None:
         self._token = token
         self._token_timestamp = int(time.time())
 
+        print(f"set token {self._token} time_stamp: {self._token_timestamp}")
+
     @property
     def auth(self) -> Tuple[str, str]:
         return self._auth
@@ -217,14 +237,21 @@ def __init__(self, factory: AuthFactory, *args: Any, **kwargs: Any) -> None:
         self._factory = factory
 
     def received_headers(self, headers: Dict[str, Any]) -> None:
-        auth_header: str = headers.get("Authorization", None)
+        auth_header: str = headers.get("authorization", None)
         if not auth_header:
             return
+        # authenticate_basic_token() returns a list. 
+        # TODO We should take the first Bearer element here
+        if isinstance(auth_header, list):
+            auth_header = auth_header[0]
+
         [auth_type, token] = auth_header.split(" ", 1)
         if auth_type == "Bearer":
             self._factory.set_token(token)
 
     def sending_headers(self) -> Dict[str, str]:
+        print("sending headers")
+
         token = self._factory.token()
         if not token:
             username, password = self._factory.auth