From d88bf83a86a72c75d9098809431d80b9172bbe96 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Fri, 29 Aug 2025 20:04:27 +0100 Subject: [PATCH] Object query service for Cassandra --- .../trustgraph/schema/services/__init__.py | 3 +- .../schema/services/objects_query.py | 28 + .../trustgraph/query/objects/__init__.py | 0 .../query/objects/cassandra/__init__.py | 0 .../query/objects/cassandra/__main__.py | 5 + .../query/objects/cassandra/service.py | 510 ++++++++++++++++++ 6 files changed, 545 insertions(+), 1 deletion(-) create mode 100644 trustgraph-base/trustgraph/schema/services/objects_query.py create mode 100644 trustgraph-flow/trustgraph/query/objects/__init__.py create mode 100644 trustgraph-flow/trustgraph/query/objects/cassandra/__init__.py create mode 100644 trustgraph-flow/trustgraph/query/objects/cassandra/__main__.py create mode 100644 trustgraph-flow/trustgraph/query/objects/cassandra/service.py diff --git a/trustgraph-base/trustgraph/schema/services/__init__.py b/trustgraph-base/trustgraph/schema/services/__init__.py index fceb0114..48711833 100644 --- a/trustgraph-base/trustgraph/schema/services/__init__.py +++ b/trustgraph-base/trustgraph/schema/services/__init__.py @@ -8,4 +8,5 @@ from .config import * from .library import * from .lookup import * from .nlp_query import * -from .structured_query import * \ No newline at end of file +from .structured_query import * +from .objects_query import * \ No newline at end of file diff --git a/trustgraph-base/trustgraph/schema/services/objects_query.py b/trustgraph-base/trustgraph/schema/services/objects_query.py new file mode 100644 index 00000000..32911b20 --- /dev/null +++ b/trustgraph-base/trustgraph/schema/services/objects_query.py @@ -0,0 +1,28 @@ +from pulsar.schema import Record, String, Map, Array + +from ..core.primitives import Error +from ..core.topic import topic + +############################################################################ + +# Objects Query Service - executes GraphQL queries against structured data + +class GraphQLError(Record): + message = String() + path = Array(String()) # Path to the field that caused the error + extensions = Map(String()) # Additional error metadata + +class ObjectsQueryRequest(Record): + user = String() # Cassandra keyspace (follows pattern from TriplesQueryRequest) + collection = String() # Data collection identifier (required for partition key) + query = String() # GraphQL query string + variables = Map(String()) # GraphQL variables + operation_name = String() # Operation to execute for multi-operation documents + +class ObjectsQueryResponse(Record): + error = Error() # System-level error (connection, timeout, etc.) + data = String() # JSON-encoded GraphQL response data + errors = Array(GraphQLError) # GraphQL field-level errors + extensions = Map(String()) # Query metadata (execution time, etc.) + +############################################################################ \ No newline at end of file diff --git a/trustgraph-flow/trustgraph/query/objects/__init__.py b/trustgraph-flow/trustgraph/query/objects/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/trustgraph-flow/trustgraph/query/objects/cassandra/__init__.py b/trustgraph-flow/trustgraph/query/objects/cassandra/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/trustgraph-flow/trustgraph/query/objects/cassandra/__main__.py b/trustgraph-flow/trustgraph/query/objects/cassandra/__main__.py new file mode 100644 index 00000000..0bec8f9d --- /dev/null +++ b/trustgraph-flow/trustgraph/query/objects/cassandra/__main__.py @@ -0,0 +1,5 @@ +#!/usr/bin/env python3 + +from . service import run + +run() \ No newline at end of file diff --git a/trustgraph-flow/trustgraph/query/objects/cassandra/service.py b/trustgraph-flow/trustgraph/query/objects/cassandra/service.py new file mode 100644 index 00000000..44d56cfd --- /dev/null +++ b/trustgraph-flow/trustgraph/query/objects/cassandra/service.py @@ -0,0 +1,510 @@ +""" +Objects query service using GraphQL. Input is a GraphQL query with variables. +Output is GraphQL response data with any errors. +""" + +import json +import logging +import asyncio +from typing import Dict, Any, Optional, List, Set +from dataclasses import dataclass, field +from cassandra.cluster import Cluster +from cassandra.auth import PlainTextAuthProvider + +import strawberry +from strawberry import Schema +from strawberry.types import Info + +from .... schema import ObjectsQueryRequest, ObjectsQueryResponse, GraphQLError +from .... schema import Error, RowSchema, Field as SchemaField +from .... base import FlowProcessor, ConsumerSpec, ProducerSpec + +# Module logger +logger = logging.getLogger(__name__) + +default_ident = "objects-query" +default_graph_host = 'localhost' + + +class Processor(FlowProcessor): + + def __init__(self, **params): + + id = params.get("id", default_ident) + + # Cassandra connection parameters + self.graph_host = params.get("graph_host", default_graph_host) + self.graph_username = params.get("graph_username", None) + self.graph_password = params.get("graph_password", None) + + # Config key for schemas + self.config_key = params.get("config_type", "schema") + + super(Processor, self).__init__( + **params | { + "id": id, + "config-type": self.config_key, + } + ) + + self.register_specification( + ConsumerSpec( + name = "request", + schema = ObjectsQueryRequest, + handler = self.on_message + ) + ) + + self.register_specification( + ProducerSpec( + name = "response", + schema = ObjectsQueryResponse, + ) + ) + + # Register config handler for schema updates + self.register_config_handler(self.on_schema_config) + + # Schema storage: name -> RowSchema + self.schemas: Dict[str, RowSchema] = {} + + # GraphQL schema + self.graphql_schema: Optional[Schema] = None + + # GraphQL types cache + self.graphql_types: Dict[str, type] = {} + + # Cassandra session + self.cluster = None + self.session = None + + # Known keyspaces and tables + self.known_keyspaces: Set[str] = set() + self.known_tables: Dict[str, Set[str]] = {} + + def connect_cassandra(self): + """Connect to Cassandra cluster""" + if self.session: + return + + try: + if self.graph_username and self.graph_password: + auth_provider = PlainTextAuthProvider( + username=self.graph_username, + password=self.graph_password + ) + self.cluster = Cluster( + contact_points=[self.graph_host], + auth_provider=auth_provider + ) + else: + self.cluster = Cluster(contact_points=[self.graph_host]) + + self.session = self.cluster.connect() + logger.info(f"Connected to Cassandra cluster at {self.graph_host}") + + except Exception as e: + logger.error(f"Failed to connect to Cassandra: {e}", exc_info=True) + raise + + def sanitize_name(self, name: str) -> str: + """Sanitize names for Cassandra compatibility""" + import re + safe_name = re.sub(r'[^a-zA-Z0-9_]', '_', name) + if safe_name and not safe_name[0].isalpha(): + safe_name = 'o_' + safe_name + return safe_name.lower() + + def sanitize_table(self, name: str) -> str: + """Sanitize table names for Cassandra compatibility""" + import re + safe_name = re.sub(r'[^a-zA-Z0-9_]', '_', name) + safe_name = 'o_' + safe_name + return safe_name.lower() + + async def on_schema_config(self, config, version): + """Handle schema configuration updates""" + logger.info(f"Loading schema configuration version {version}") + + # Clear existing schemas + self.schemas = {} + self.graphql_types = {} + + # Check if our config type exists + if self.config_key not in config: + logger.warning(f"No '{self.config_key}' type in configuration") + return + + # Get the schemas dictionary for our type + schemas_config = config[self.config_key] + + # Process each schema in the schemas config + for schema_name, schema_json in schemas_config.items(): + try: + # Parse the JSON schema definition + schema_def = json.loads(schema_json) + + # Create Field objects + fields = [] + for field_def in schema_def.get("fields", []): + field = SchemaField( + name=field_def["name"], + type=field_def["type"], + size=field_def.get("size", 0), + primary=field_def.get("primary_key", False), + description=field_def.get("description", ""), + required=field_def.get("required", False), + enum_values=field_def.get("enum", []), + indexed=field_def.get("indexed", False) + ) + fields.append(field) + + # Create RowSchema + row_schema = RowSchema( + name=schema_def.get("name", schema_name), + description=schema_def.get("description", ""), + fields=fields + ) + + self.schemas[schema_name] = row_schema + logger.info(f"Loaded schema: {schema_name} with {len(fields)} fields") + + except Exception as e: + logger.error(f"Failed to parse schema {schema_name}: {e}", exc_info=True) + + logger.info(f"Schema configuration loaded: {len(self.schemas)} schemas") + + # Regenerate GraphQL schema + self.generate_graphql_schema() + + def get_python_type(self, field_type: str): + """Convert schema field type to Python type for GraphQL""" + type_mapping = { + "string": str, + "integer": int, + "float": float, + "boolean": bool, + "timestamp": str, # Use string for timestamps in GraphQL + "date": str, + "time": str, + "uuid": str + } + return type_mapping.get(field_type, str) + + def create_graphql_type(self, schema_name: str, row_schema: RowSchema) -> type: + """Create a GraphQL type from a RowSchema""" + + # Create annotations for the GraphQL type + annotations = {} + defaults = {} + + for field in row_schema.fields: + python_type = self.get_python_type(field.type) + + # Make field optional if not required + if not field.required and not field.primary: + annotations[field.name] = Optional[python_type] + defaults[field.name] = None + else: + annotations[field.name] = python_type + + # Create the class dynamically + type_name = f"{schema_name.capitalize()}Type" + graphql_class = type( + type_name, + (), + { + "__annotations__": annotations, + **defaults + } + ) + + # Apply strawberry decorator + return strawberry.type(graphql_class) + + def generate_graphql_schema(self): + """Generate GraphQL schema from loaded schemas""" + if not self.schemas: + logger.warning("No schemas loaded, cannot generate GraphQL schema") + self.graphql_schema = None + return + + # Create GraphQL types for each schema + for schema_name, row_schema in self.schemas.items(): + graphql_type = self.create_graphql_type(schema_name, row_schema) + self.graphql_types[schema_name] = graphql_type + + # Create the Query class with resolvers + query_dict = {'__annotations__': {}} + + for schema_name, row_schema in self.schemas.items(): + graphql_type = self.graphql_types[schema_name] + + # Create resolver function for this schema + def make_resolver(schema_name, row_schema, graphql_type): + async def resolver( + self_query, + info: Info, + collection: str, + limit: Optional[int] = 100, + **filters + ) -> List[Any]: + # Get the processor instance from context + processor = info.context["processor"] + user = info.context["user"] + + # Query Cassandra + results = await processor.query_cassandra( + user, collection, schema_name, row_schema, + filters, limit + ) + + # Convert to GraphQL types + graphql_results = [] + for row in results: + graphql_obj = graphql_type(**row) + graphql_results.append(graphql_obj) + + return graphql_results + + return resolver + + # Add resolver to query + resolver_name = f"{schema_name}_objects" + resolver_func = make_resolver(schema_name, row_schema, graphql_type) + + # Create field with proper annotations + query_dict[resolver_name] = strawberry.field(resolver_func) + query_dict['__annotations__'][resolver_name] = List[graphql_type] + + # Create the Query class + Query = type('Query', (), query_dict) + Query = strawberry.type(Query) + + # Create the schema + self.graphql_schema = strawberry.Schema(query=Query) + logger.info(f"Generated GraphQL schema with {len(self.schemas)} types") + + async def query_cassandra( + self, + user: str, + collection: str, + schema_name: str, + row_schema: RowSchema, + filters: Dict[str, Any], + limit: int + ) -> List[Dict[str, Any]]: + """Execute a query against Cassandra""" + + # Connect if needed + self.connect_cassandra() + + # Build the query + keyspace = self.sanitize_name(user) + table = self.sanitize_table(schema_name) + + # Start with basic SELECT + query = f"SELECT * FROM {keyspace}.{table}" + + # Add WHERE clauses + where_clauses = [f"collection = %s"] + params = [collection] + + # Add filters for indexed or primary key fields + for field_name, value in filters.items(): + if value is not None: + # Find the field in schema + schema_field = None + for f in row_schema.fields: + if f.name == field_name: + schema_field = f + break + + if schema_field and (schema_field.indexed or schema_field.primary): + safe_field = self.sanitize_name(field_name) + where_clauses.append(f"{safe_field} = %s") + params.append(value) + + if where_clauses: + query += " WHERE " + " AND ".join(where_clauses) + + # Add limit + if limit: + query += f" LIMIT {limit}" + + # Execute query + try: + result = self.session.execute(query, params) + + # Convert rows to dicts + results = [] + for row in result: + row_dict = {} + for field in row_schema.fields: + safe_field = self.sanitize_name(field.name) + if hasattr(row, safe_field): + value = getattr(row, safe_field) + # Use original field name in result + row_dict[field.name] = value + results.append(row_dict) + + return results + + except Exception as e: + logger.error(f"Failed to query Cassandra: {e}", exc_info=True) + raise + + async def execute_graphql_query( + self, + query: str, + variables: Dict[str, Any], + operation_name: Optional[str], + user: str, + collection: str + ) -> Dict[str, Any]: + """Execute a GraphQL query""" + + if not self.graphql_schema: + raise RuntimeError("No GraphQL schema available - no schemas loaded") + + # Create context for the query + context = { + "processor": self, + "user": user, + "collection": collection + } + + # Execute the query + result = await self.graphql_schema.execute( + query, + variable_values=variables, + operation_name=operation_name, + context_value=context + ) + + # Build response + response = {} + + if result.data: + response["data"] = result.data + else: + response["data"] = None + + if result.errors: + response["errors"] = [ + { + "message": str(error), + "path": getattr(error, "path", []), + "extensions": getattr(error, "extensions", {}) + } + for error in result.errors + ] + else: + response["errors"] = [] + + # Add extensions if any + if hasattr(result, "extensions") and result.extensions: + response["extensions"] = result.extensions + + return response + + async def on_message(self, msg, consumer, flow): + """Handle incoming query request""" + + try: + request = msg.value() + + # Sender-produced ID + id = msg.properties()["id"] + + logger.debug(f"Handling objects query request {id}...") + + # Execute GraphQL query + result = await self.execute_graphql_query( + query=request.query, + variables=dict(request.variables) if request.variables else {}, + operation_name=request.operation_name, + user=request.user, + collection=request.collection + ) + + # Create response + graphql_errors = [] + if "errors" in result and result["errors"]: + for err in result["errors"]: + graphql_error = GraphQLError( + message=err.get("message", ""), + path=err.get("path", []), + extensions=err.get("extensions", {}) + ) + graphql_errors.append(graphql_error) + + response = ObjectsQueryResponse( + error=None, + data=json.dumps(result.get("data")) if result.get("data") else "null", + errors=graphql_errors, + extensions=result.get("extensions", {}) + ) + + logger.debug("Sending objects query response...") + await flow("response").send(response, properties={"id": id}) + + logger.debug("Objects query request completed") + + except Exception as e: + + logger.error(f"Exception in objects query service: {e}", exc_info=True) + + logger.info("Sending error response...") + + response = ObjectsQueryResponse( + error = Error( + type = "objects-query-error", + message = str(e), + ), + data = None, + errors = [], + extensions = {} + ) + + await flow("response").send(response, properties={"id": id}) + + def close(self): + """Clean up Cassandra connections""" + if self.cluster: + self.cluster.shutdown() + logger.info("Closed Cassandra connection") + + @staticmethod + def add_args(parser): + """Add command-line arguments""" + + FlowProcessor.add_args(parser) + + parser.add_argument( + '-g', '--graph-host', + default=default_graph_host, + help=f'Cassandra host (default: {default_graph_host})' + ) + + parser.add_argument( + '--graph-username', + default=None, + help='Cassandra username' + ) + + parser.add_argument( + '--graph-password', + default=None, + help='Cassandra password' + ) + + parser.add_argument( + '--config-type', + default='schema', + help='Configuration type prefix for schemas (default: schema)' + ) + +def run(): + """Entry point for objects-query-graphql-cassandra command""" + Processor.launch(default_ident, __doc__) \ No newline at end of file