Add snowgraph support

mcp_server/pyproject.toml

@@ -16,6 +16,7 @@ classifiers = [
 dependencies = [
     "graphdatascience>=1.16",
     "mcp[cli]>=1.11.0",
+    "snowflake-snowpark-python==1.38.0",
 ]
 
 [project.urls]

mcp_server/src/mcp_server_neo4j_gds/centrality_algorithm_handlers.py

@@ -187,32 +187,82 @@ def execute(self, arguments: Dict[str, Any]) -> Any:
 
 class DegreeCentralityHandler(AlgorithmHandler):
     def degree_centrality(self, **kwargs):
-        with projected_graph(self.gds) as G:
-            params = {
-                k: v
-                for k, v in kwargs.items()
-                if v is not None and k not in ["nodes", "nodeIdentifierProperty"]
+        # with projected_graph(self.gds) as G:
+        #     params = {
+        #         k: v
+        #         for k, v in kwargs.items()
+        #         if v is not None and k not in ["nodes", "nodeIdentifierProperty"]
+        #     }
+        #     logger.info(f"Degree centrality parameters: {params}")
+        #     centrality = self.gds.degree.stream(G, **params)
+
+        # # Add node names to the results if nodeIdentifierProperty is provided
+        # node_identifier_property = kwargs.get("nodeIdentifierProperty")
+        # translate_ids_to_identifiers(self.gds, node_identifier_property, centrality)
+
+        # # Filter results by node names if provided
+        # node_names = kwargs.get("nodes", None)
+        # centrality = filter_identifiers(
+        #     self.gds, node_identifier_property, node_names, centrality
+        # )
+        orientation = kwargs.get("orientation")
+        if orientation is None:
+            relationships_dict = """
+            {
+                'RELATIONSHIPS': {
+                    'sourceTable': 'NODES',
+                    'targetTable': 'NODES',
+                }
             }
-            logger.info(f"Degree centrality parameters: {params}")
-            centrality = self.gds.degree.stream(G, **params)
-
-        # Add node names to the results if nodeIdentifierProperty is provided
-        node_identifier_property = kwargs.get("nodeIdentifierProperty")
-        translate_ids_to_identifiers(self.gds, node_identifier_property, centrality)
-
-        # Filter results by node names if provided
-        node_names = kwargs.get("nodes", None)
-        centrality = filter_identifiers(
-            self.gds, node_identifier_property, node_names, centrality
-        )
-
-        return centrality
+            """
+        else:
+            relationships_dict = f"""
+            {{
+                'RELATIONSHIPS': {{
+                    'sourceTable': 'NODES',
+                    'targetTable': 'NODES',
+                    'orientation': '{orientation}'
+                }}
+            }}
+            """
+
+        relationshipWeightProperty = kwargs.get("relationshipWeightProperty")
+        if relationshipWeightProperty is None:
+            compute_dict = """
+            { }
+            """
+        else:
+            compute_dict = f"""
+            {{
+                'relationshipWeightProperty': '{relationshipWeightProperty}'
+            }}
+            """
+
+        res = self.gds.sql(
+                f"""
+                CALL Neo4j_Graph_Analytics.graph.degree('CPU_X64_XS', {{
+                    'defaultTablePrefix': 'EXAMPLE_DB.PUBLIC',
+                    'project': {{
+                        'nodeTables': [ 'NODES' ],
+                        'relationshipTables': {relationships_dict}
+                    }},
+                    'compute': {compute_dict},
+                    'write': [{{
+                        'nodeLabel': 'NODES',
+                        'outputTable': 'NODES_DEGREE_CENTRALITY'
+                    }}]
+                }});
+                """
+        ).collect()
+
+        logger.info(f"Degree centrality execution: {res}")
+        output_table = self.gds.table("EXAMPLE_DB.PUBLIC.NODES_DEGREE_CENTRALITY")
+        return output_table.to_pandas()
 
     def execute(self, arguments: Dict[str, Any]) -> Any:
         return self.degree_centrality(
-            nodes=arguments.get("nodes"),
-            nodeIdentifierProperty=arguments.get("nodeIdentifierProperty"),
             orientation=arguments.get("orientation"),
+            relationshipWeightProperty=arguments.get("relationshipWeightProperty"),
         )
 
 

mcp_server/src/mcp_server_neo4j_gds/centrality_algorithm_specs.py

@@ -177,19 +177,14 @@
         inputSchema={
             "type": "object",
             "properties": {
-                "nodes": {
-                    "type": "array",
-                    "items": {"type": "string"},
-                    "description": "List of node names to filter degree centrality results for.",
-                },
-                "nodeIdentifierProperty": {
-                    "type": "string",
-                    "description": "Property name to use for identifying nodes (e.g., 'name', 'Name', 'title'). Use get_node_properties_keys to find available properties.",
-                },
                 "orientation": {
                     "type": "string",
                     "description": "The orientation used to compute node degrees. Supported orientations are NATURAL (for out-degree), REVERSE (for in-degree) and UNDIRECTED (for both in-degree and out-degree) ",
                 },
+                "relationshipWeightProperty": {
+                    "type": "string",
+                    "description": "Property of the relationship to use for weighting. If not specified, all relationships are treated equally.",
+                },
             },
             "required": [],
         },

mcp_server/src/mcp_server_neo4j_gds/server.py

@@ -8,7 +8,7 @@
 import pandas as pd
 import json
 from graphdatascience import GraphDataScience
-
+from snowflake.snowpark import Session
 from .similarity_algorithm_specs import similarity_tool_definitions
 from .centrality_algorithm_specs import centrality_tool_definitions
 from .community_algorithm_specs import community_tool_definitions
@@ -56,7 +56,12 @@ async def main(db_url: str, username: str, password: str, database: str = None):
                 db_url, auth=(username, password), aura_ds=False, database=database
             )
         else:
-            gds = GraphDataScience(db_url, auth=(username, password), aura_ds=False)
+            if db_url:
+                gds = GraphDataScience(db_url, auth=(username, password), aura_ds=False)
+            else:
+                gds = Session.builder.config("connection_name", "snowflake-gds-mcp").create()
+                print(gds.sql("SELECT 1;").collect())
+                logger.info("Successfully connected to Snowflake database")
         logger.info("Successfully connected to Neo4j database")
     except Exception as e:
         logger.error(f"Failed to connect to Neo4j database: {e}")