Update notebook

Author: FlorentinD

examples/dev/gds-sessions.ipynb

@@ -70,18 +70,21 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from graphdatascience.gds_session.dbms_connection_info import DbmsConnectionInfo\n",
+    "from graphdatascience.gds_session import DbmsConnectionInfo\n",
     "import os\n",
     "\n",
     "from neo4j import GraphDatabase\n",
     "\n",
-    "# We need to tell the GDS client that we are working with a development environment.\n",
-    "# This does not need to be set in production.\n",
-    "os.environ[\"AURA_ENV\"] = \"devstrawberryfield\"\n",
+    "# We can tell the GDS client that we are working with a development environment.\n",
+    "# os.environ[\"AURA_ENV\"] = \"devstrawberryfield\"\n",
     "\n",
-    "db_connection_info = DbmsConnectionInfo(\n",
-    "    f\"neo4j+s://{db_id}-{os.environ['AURA_ENV']}.databases.neo4j-dev.io\", \"neo4j\", db_password\n",
+    "uri = (\n",
+    "    f\"neo4j+s://{db_id}-{os.environ['AURA_ENV']}.databases.neo4j-dev.io\"\n",
+    "    if os.environ.get(\"AURA_ENV\")\n",
+    "    else f\"neo4j+s://{db_id}.databases.neo4j.io\"\n",
     ")\n",
+    "\n",
+    "db_connection_info = DbmsConnectionInfo(uri, \"neo4j\", db_password)\n",
     "# start a standard Neo4j Python Driver to connect to the AuraDB instance\n",
     "driver = GraphDatabase.driver(db_connection_info.uri, auth=db_connection_info.auth())\n",
     "\n",
@@ -171,9 +174,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Let's call this function and see what it returns\n",
+    "# Let's call this function and verify the arrow server is enabled\n",
     "with driver.session() as session:\n",
-    "    display(session.run(\"CALL internal.arrow.status\").to_df())"
+    "    arrow_status = session.run(\"CALL internal.arrow.status\").to_df()\n",
+    "    print(arrow_status)\n",
+    "    assert arrow_status[\"enabled\"][0] == True, \"Arrow server on the db needs to be enabled\""
    ]
   },
   {
@@ -244,7 +249,7 @@
    "outputs": [],
    "source": [
     "# The new stuff!\n",
-    "from graphdatascience.gds_session.gds_sessions import GdsSessions, AuraAPICredentials\n",
+    "from graphdatascience.gds_session import GdsSessions, AuraAPICredentials\n",
     "\n",
     "# Create a new AuraSessions object\n",
     "sessions = GdsSessions(ds_connection=AuraAPICredentials(CLIENT_ID, CLIENT_SECRET))"
@@ -311,6 +316,7 @@
     "    session_name=\"pagerank-compute\",\n",
     "    size=SessionSizes.by_memory().DEFAULT,\n",
     "    db_connection=db_connection_info,\n",
+    "    region=\"europe-west1\",\n",
     ")\n",
     "\n",
     "# If we had already created a session and want to reconnect to it, the same code is used.\n",
@@ -542,8 +548,22 @@
   }
  ],
  "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
   "language_info": {
-   "name": "python"
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.7"
   }
  },
  "nbformat": 4,