From 65b3a6ee12147ac968c70ab6571f10268f779b67 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Mon, 16 Feb 2026 11:34:40 +0000 Subject: [PATCH] Graph implementation --- tests/pytest.ini | 3 +- .../test_direct/test_entity_centric_kg.py | 643 ++++++++++++++++++ .../test_triples_cassandra_query.py | 306 +++++---- .../test_triples_cassandra_storage.py | 209 +++--- .../trustgraph/direct/cassandra_kg.py | 587 ++++++++++++++++ .../query/triples/cassandra/service.py | 79 ++- .../storage/triples/cassandra/write.py | 77 ++- 7 files changed, 1679 insertions(+), 225 deletions(-) create mode 100644 tests/unit/test_direct/test_entity_centric_kg.py diff --git a/tests/pytest.ini b/tests/pytest.ini index b763299c..b032a9d4 100644 --- a/tests/pytest.ini +++ b/tests/pytest.ini @@ -19,4 +19,5 @@ markers = integration: marks tests as integration tests unit: marks tests as unit tests contract: marks tests as contract tests (service interface validation) - vertexai: marks tests as vertex ai specific tests \ No newline at end of file + vertexai: marks tests as vertex ai specific tests + asyncio: marks tests that use asyncio \ No newline at end of file diff --git a/tests/unit/test_direct/test_entity_centric_kg.py b/tests/unit/test_direct/test_entity_centric_kg.py new file mode 100644 index 00000000..af77dab1 --- /dev/null +++ b/tests/unit/test_direct/test_entity_centric_kg.py @@ -0,0 +1,643 @@ +""" +Unit tests for EntityCentricKnowledgeGraph class + +Tests the entity-centric knowledge graph implementation without requiring +an actual Cassandra connection. Uses mocking to verify correct behavior. +""" + +import pytest +from unittest.mock import MagicMock, patch, call +import os + + +class TestEntityCentricKnowledgeGraph: + """Test cases for EntityCentricKnowledgeGraph""" + + @pytest.fixture + def mock_cluster(self): + """Create a mock Cassandra cluster""" + with patch('trustgraph.direct.cassandra_kg.Cluster') as mock_cluster_cls: + mock_cluster = MagicMock() + mock_session = MagicMock() + mock_cluster.connect.return_value = mock_session + mock_cluster_cls.return_value = mock_cluster + yield mock_cluster_cls, mock_cluster, mock_session + + @pytest.fixture + def entity_kg(self, mock_cluster): + """Create an EntityCentricKnowledgeGraph instance with mocked Cassandra""" + from trustgraph.direct.cassandra_kg import EntityCentricKnowledgeGraph + mock_cluster_cls, mock_cluster, mock_session = mock_cluster + + # Create instance + kg = EntityCentricKnowledgeGraph(hosts=['localhost'], keyspace='test_keyspace') + return kg, mock_session + + def test_init_creates_entity_centric_schema(self, mock_cluster): + """Test that initialization creates the 2-table entity-centric schema""" + from trustgraph.direct.cassandra_kg import EntityCentricKnowledgeGraph + mock_cluster_cls, mock_cluster, mock_session = mock_cluster + + kg = EntityCentricKnowledgeGraph(hosts=['localhost'], keyspace='test_keyspace') + + # Verify schema tables were created + execute_calls = mock_session.execute.call_args_list + executed_statements = [str(c) for c in execute_calls] + + # Check for keyspace creation + keyspace_created = any('create keyspace' in str(c).lower() for c in execute_calls) + assert keyspace_created + + # Check for quads_by_entity table + entity_table_created = any('quads_by_entity' in str(c) for c in execute_calls) + assert entity_table_created + + # Check for quads_by_collection table + collection_table_created = any('quads_by_collection' in str(c) for c in execute_calls) + assert collection_table_created + + # Check for collection_metadata table + metadata_table_created = any('collection_metadata' in str(c) for c in execute_calls) + assert metadata_table_created + + def test_prepare_statements_initialized(self, entity_kg): + """Test that prepared statements are initialized""" + kg, mock_session = entity_kg + + # Verify prepare was called for various statements + assert mock_session.prepare.called + prepare_calls = mock_session.prepare.call_args_list + + # Check that key prepared statements exist + prepared_queries = [str(c) for c in prepare_calls] + + # Insert statements + insert_entity_stmt = any('INSERT INTO' in str(c) and 'quads_by_entity' in str(c) + for c in prepare_calls) + assert insert_entity_stmt + + insert_collection_stmt = any('INSERT INTO' in str(c) and 'quads_by_collection' in str(c) + for c in prepare_calls) + assert insert_collection_stmt + + def test_insert_uri_object_creates_4_entity_rows(self, entity_kg): + """Test that inserting a quad with URI object creates 4 entity rows""" + kg, mock_session = entity_kg + + # Reset mocks to track only insert-related calls + mock_session.reset_mock() + + kg.insert( + collection='test_collection', + s='http://example.org/Alice', + p='http://example.org/knows', + o='http://example.org/Bob', + g='http://example.org/graph1', + otype='U' + ) + + # Verify batch was executed + mock_session.execute.assert_called() + + def test_insert_literal_object_creates_3_entity_rows(self, entity_kg): + """Test that inserting a quad with literal object creates 3 entity rows""" + kg, mock_session = entity_kg + + mock_session.reset_mock() + + kg.insert( + collection='test_collection', + s='http://example.org/Alice', + p='http://www.w3.org/2000/01/rdf-schema#label', + o='Alice Smith', + g=None, + otype='L', + dtype='xsd:string', + lang='en' + ) + + # Verify batch was executed + mock_session.execute.assert_called() + + def test_insert_default_graph(self, entity_kg): + """Test that None graph is stored as empty string""" + kg, mock_session = entity_kg + + mock_session.reset_mock() + + kg.insert( + collection='test_collection', + s='http://example.org/Alice', + p='http://example.org/knows', + o='http://example.org/Bob', + g=None, + otype='U' + ) + + mock_session.execute.assert_called() + + def test_insert_auto_detects_otype(self, entity_kg): + """Test that otype is auto-detected when not provided""" + kg, mock_session = entity_kg + + mock_session.reset_mock() + + # URI should be auto-detected + kg.insert( + collection='test_collection', + s='http://example.org/Alice', + p='http://example.org/knows', + o='http://example.org/Bob' + ) + mock_session.execute.assert_called() + + mock_session.reset_mock() + + # Literal should be auto-detected + kg.insert( + collection='test_collection', + s='http://example.org/Alice', + p='http://example.org/name', + o='Alice' + ) + mock_session.execute.assert_called() + + def test_get_s_returns_quads_for_subject(self, entity_kg): + """Test get_s queries by subject""" + kg, mock_session = entity_kg + + # Mock the query result + mock_result = [ + MagicMock(p='http://example.org/knows', o='http://example.org/Bob', + d='', otype='U', dtype='', lang='', s='http://example.org/Alice') + ] + mock_session.execute.return_value = mock_result + + results = kg.get_s('test_collection', 'http://example.org/Alice') + + # Verify query was executed + mock_session.execute.assert_called() + + # Results should be QuadResult objects + assert len(results) == 1 + assert results[0].s == 'http://example.org/Alice' + assert results[0].p == 'http://example.org/knows' + assert results[0].o == 'http://example.org/Bob' + + def test_get_p_returns_quads_for_predicate(self, entity_kg): + """Test get_p queries by predicate""" + kg, mock_session = entity_kg + + mock_result = [ + MagicMock(s='http://example.org/Alice', o='http://example.org/Bob', + d='', otype='U', dtype='', lang='', p='http://example.org/knows') + ] + mock_session.execute.return_value = mock_result + + results = kg.get_p('test_collection', 'http://example.org/knows') + + mock_session.execute.assert_called() + assert len(results) == 1 + + def test_get_o_returns_quads_for_object(self, entity_kg): + """Test get_o queries by object""" + kg, mock_session = entity_kg + + mock_result = [ + MagicMock(s='http://example.org/Alice', p='http://example.org/knows', + d='', otype='U', dtype='', lang='', o='http://example.org/Bob') + ] + mock_session.execute.return_value = mock_result + + results = kg.get_o('test_collection', 'http://example.org/Bob') + + mock_session.execute.assert_called() + assert len(results) == 1 + + def test_get_sp_returns_quads_for_subject_predicate(self, entity_kg): + """Test get_sp queries by subject and predicate""" + kg, mock_session = entity_kg + + mock_result = [ + MagicMock(o='http://example.org/Bob', d='', otype='U', dtype='', lang='') + ] + mock_session.execute.return_value = mock_result + + results = kg.get_sp('test_collection', 'http://example.org/Alice', + 'http://example.org/knows') + + mock_session.execute.assert_called() + assert len(results) == 1 + + def test_get_po_returns_quads_for_predicate_object(self, entity_kg): + """Test get_po queries by predicate and object""" + kg, mock_session = entity_kg + + mock_result = [ + MagicMock(s='http://example.org/Alice', d='', otype='U', dtype='', lang='', + o='http://example.org/Bob') + ] + mock_session.execute.return_value = mock_result + + results = kg.get_po('test_collection', 'http://example.org/knows', + 'http://example.org/Bob') + + mock_session.execute.assert_called() + assert len(results) == 1 + + def test_get_os_returns_quads_for_object_subject(self, entity_kg): + """Test get_os queries by object and subject""" + kg, mock_session = entity_kg + + mock_result = [ + MagicMock(p='http://example.org/knows', d='', otype='U', dtype='', lang='', + s='http://example.org/Alice', o='http://example.org/Bob') + ] + mock_session.execute.return_value = mock_result + + results = kg.get_os('test_collection', 'http://example.org/Bob', + 'http://example.org/Alice') + + mock_session.execute.assert_called() + assert len(results) == 1 + + def test_get_spo_returns_quads_for_subject_predicate_object(self, entity_kg): + """Test get_spo queries by subject, predicate, and object""" + kg, mock_session = entity_kg + + mock_result = [ + MagicMock(d='', otype='U', dtype='', lang='', + o='http://example.org/Bob') + ] + mock_session.execute.return_value = mock_result + + results = kg.get_spo('test_collection', 'http://example.org/Alice', + 'http://example.org/knows', 'http://example.org/Bob') + + mock_session.execute.assert_called() + assert len(results) == 1 + + def test_get_g_returns_quads_for_graph(self, entity_kg): + """Test get_g queries by graph""" + kg, mock_session = entity_kg + + mock_result = [ + MagicMock(s='http://example.org/Alice', p='http://example.org/knows', + o='http://example.org/Bob', otype='U', dtype='', lang='') + ] + mock_session.execute.return_value = mock_result + + results = kg.get_g('test_collection', 'http://example.org/graph1') + + mock_session.execute.assert_called() + + def test_get_all_returns_all_quads_in_collection(self, entity_kg): + """Test get_all returns all quads""" + kg, mock_session = entity_kg + + mock_result = [ + MagicMock(d='', s='http://example.org/Alice', p='http://example.org/knows', + o='http://example.org/Bob', otype='U', dtype='', lang='') + ] + mock_session.execute.return_value = mock_result + + results = kg.get_all('test_collection') + + mock_session.execute.assert_called() + + def test_graph_wildcard_returns_all_graphs(self, entity_kg): + """Test that g='*' returns quads from all graphs""" + from trustgraph.direct.cassandra_kg import GRAPH_WILDCARD + kg, mock_session = entity_kg + + mock_result = [ + MagicMock(p='http://example.org/knows', d='http://example.org/graph1', + otype='U', dtype='', lang='', s='http://example.org/Alice', + o='http://example.org/Bob'), + MagicMock(p='http://example.org/knows', d='http://example.org/graph2', + otype='U', dtype='', lang='', s='http://example.org/Alice', + o='http://example.org/Charlie') + ] + mock_session.execute.return_value = mock_result + + results = kg.get_s('test_collection', 'http://example.org/Alice', g=GRAPH_WILDCARD) + + # Should return quads from both graphs + assert len(results) == 2 + + def test_specific_graph_filters_results(self, entity_kg): + """Test that specifying a graph filters results""" + kg, mock_session = entity_kg + + mock_result = [ + MagicMock(p='http://example.org/knows', d='http://example.org/graph1', + otype='U', dtype='', lang='', s='http://example.org/Alice', + o='http://example.org/Bob'), + MagicMock(p='http://example.org/knows', d='http://example.org/graph2', + otype='U', dtype='', lang='', s='http://example.org/Alice', + o='http://example.org/Charlie') + ] + mock_session.execute.return_value = mock_result + + results = kg.get_s('test_collection', 'http://example.org/Alice', + g='http://example.org/graph1') + + # Should only return quads from graph1 + assert len(results) == 1 + assert results[0].g == 'http://example.org/graph1' + + def test_collection_exists_returns_true_when_exists(self, entity_kg): + """Test collection_exists returns True for existing collection""" + kg, mock_session = entity_kg + + mock_result = [MagicMock(collection='test_collection')] + mock_session.execute.return_value = mock_result + + exists = kg.collection_exists('test_collection') + + assert exists is True + + def test_collection_exists_returns_false_when_not_exists(self, entity_kg): + """Test collection_exists returns False for non-existing collection""" + kg, mock_session = entity_kg + + mock_session.execute.return_value = [] + + exists = kg.collection_exists('nonexistent_collection') + + assert exists is False + + def test_create_collection_inserts_metadata(self, entity_kg): + """Test create_collection inserts metadata row""" + kg, mock_session = entity_kg + + mock_session.reset_mock() + kg.create_collection('test_collection') + + # Verify INSERT was executed for collection_metadata + mock_session.execute.assert_called() + + def test_delete_collection_removes_all_data(self, entity_kg): + """Test delete_collection removes entity partitions and collection rows""" + kg, mock_session = entity_kg + + # Mock reading quads from collection + mock_quads = [ + MagicMock(d='', s='http://example.org/Alice', p='http://example.org/knows', + o='http://example.org/Bob', otype='U') + ] + mock_session.execute.return_value = mock_quads + + mock_session.reset_mock() + kg.delete_collection('test_collection') + + # Verify delete operations were executed + assert mock_session.execute.called + + def test_close_shuts_down_connections(self, entity_kg): + """Test close shuts down session and cluster""" + kg, mock_session = entity_kg + + kg.close() + + mock_session.shutdown.assert_called_once() + kg.cluster.shutdown.assert_called_once() + + +class TestQuadResult: + """Test cases for QuadResult class""" + + def test_quad_result_stores_all_fields(self): + """Test QuadResult stores all quad fields""" + from trustgraph.direct.cassandra_kg import QuadResult + + result = QuadResult( + s='http://example.org/Alice', + p='http://example.org/knows', + o='http://example.org/Bob', + g='http://example.org/graph1', + otype='U', + dtype='', + lang='' + ) + + assert result.s == 'http://example.org/Alice' + assert result.p == 'http://example.org/knows' + assert result.o == 'http://example.org/Bob' + assert result.g == 'http://example.org/graph1' + assert result.otype == 'U' + assert result.dtype == '' + assert result.lang == '' + + def test_quad_result_defaults(self): + """Test QuadResult default values""" + from trustgraph.direct.cassandra_kg import QuadResult + + result = QuadResult( + s='http://example.org/s', + p='http://example.org/p', + o='literal value', + g='' + ) + + assert result.otype == 'U' # Default otype + assert result.dtype == '' + assert result.lang == '' + + def test_quad_result_with_literal_metadata(self): + """Test QuadResult with literal metadata""" + from trustgraph.direct.cassandra_kg import QuadResult + + result = QuadResult( + s='http://example.org/Alice', + p='http://www.w3.org/2000/01/rdf-schema#label', + o='Alice Smith', + g='', + otype='L', + dtype='xsd:string', + lang='en' + ) + + assert result.otype == 'L' + assert result.dtype == 'xsd:string' + assert result.lang == 'en' + + +class TestFactoryFunction: + """Test cases for get_knowledge_graph_class factory function""" + + def test_factory_returns_legacy_by_default(self): + """Test factory returns KnowledgeGraph when env var not set""" + from trustgraph.direct.cassandra_kg import get_knowledge_graph_class, KnowledgeGraph + + with patch.dict(os.environ, {}, clear=True): + KGClass = get_knowledge_graph_class() + assert KGClass == KnowledgeGraph + + def test_factory_returns_legacy_when_false(self): + """Test factory returns KnowledgeGraph when env var is false""" + from trustgraph.direct.cassandra_kg import get_knowledge_graph_class, KnowledgeGraph + + with patch.dict(os.environ, {'CASSANDRA_ENTITY_CENTRIC': 'false'}): + KGClass = get_knowledge_graph_class() + assert KGClass == KnowledgeGraph + + def test_factory_returns_entity_centric_when_true(self): + """Test factory returns EntityCentricKnowledgeGraph when env var is true""" + from trustgraph.direct.cassandra_kg import ( + get_knowledge_graph_class, EntityCentricKnowledgeGraph + ) + + with patch.dict(os.environ, {'CASSANDRA_ENTITY_CENTRIC': 'true'}): + KGClass = get_knowledge_graph_class() + assert KGClass == EntityCentricKnowledgeGraph + + def test_factory_case_insensitive(self): + """Test factory handles case insensitive values""" + from trustgraph.direct.cassandra_kg import ( + get_knowledge_graph_class, EntityCentricKnowledgeGraph + ) + + with patch.dict(os.environ, {'CASSANDRA_ENTITY_CENTRIC': 'TRUE'}): + KGClass = get_knowledge_graph_class() + assert KGClass == EntityCentricKnowledgeGraph + + with patch.dict(os.environ, {'CASSANDRA_ENTITY_CENTRIC': 'True'}): + KGClass = get_knowledge_graph_class() + assert KGClass == EntityCentricKnowledgeGraph + + +class TestWriteHelperFunctions: + """Test cases for helper functions in write.py""" + + def test_get_term_otype_for_iri(self): + """Test get_term_otype returns 'U' for IRI terms""" + from trustgraph.storage.triples.cassandra.write import get_term_otype + from trustgraph.schema import Term, IRI + + term = Term(type=IRI, iri='http://example.org/Alice') + assert get_term_otype(term) == 'U' + + def test_get_term_otype_for_literal(self): + """Test get_term_otype returns 'L' for LITERAL terms""" + from trustgraph.storage.triples.cassandra.write import get_term_otype + from trustgraph.schema import Term, LITERAL + + term = Term(type=LITERAL, value='Alice Smith') + assert get_term_otype(term) == 'L' + + def test_get_term_otype_for_blank(self): + """Test get_term_otype returns 'U' for BLANK terms""" + from trustgraph.storage.triples.cassandra.write import get_term_otype + from trustgraph.schema import Term, BLANK + + term = Term(type=BLANK, id='_:b1') + assert get_term_otype(term) == 'U' + + def test_get_term_otype_for_triple(self): + """Test get_term_otype returns 'T' for TRIPLE terms""" + from trustgraph.storage.triples.cassandra.write import get_term_otype + from trustgraph.schema import Term, TRIPLE + + term = Term(type=TRIPLE) + assert get_term_otype(term) == 'T' + + def test_get_term_otype_for_none(self): + """Test get_term_otype returns 'U' for None""" + from trustgraph.storage.triples.cassandra.write import get_term_otype + + assert get_term_otype(None) == 'U' + + def test_get_term_dtype_for_literal(self): + """Test get_term_dtype extracts datatype from LITERAL""" + from trustgraph.storage.triples.cassandra.write import get_term_dtype + from trustgraph.schema import Term, LITERAL + + term = Term(type=LITERAL, value='42', datatype='xsd:integer') + assert get_term_dtype(term) == 'xsd:integer' + + def test_get_term_dtype_for_non_literal(self): + """Test get_term_dtype returns empty string for non-LITERAL""" + from trustgraph.storage.triples.cassandra.write import get_term_dtype + from trustgraph.schema import Term, IRI + + term = Term(type=IRI, iri='http://example.org/Alice') + assert get_term_dtype(term) == '' + + def test_get_term_dtype_for_none(self): + """Test get_term_dtype returns empty string for None""" + from trustgraph.storage.triples.cassandra.write import get_term_dtype + + assert get_term_dtype(None) == '' + + def test_get_term_lang_for_literal(self): + """Test get_term_lang extracts language from LITERAL""" + from trustgraph.storage.triples.cassandra.write import get_term_lang + from trustgraph.schema import Term, LITERAL + + term = Term(type=LITERAL, value='Alice Smith', language='en') + assert get_term_lang(term) == 'en' + + def test_get_term_lang_for_non_literal(self): + """Test get_term_lang returns empty string for non-LITERAL""" + from trustgraph.storage.triples.cassandra.write import get_term_lang + from trustgraph.schema import Term, IRI + + term = Term(type=IRI, iri='http://example.org/Alice') + assert get_term_lang(term) == '' + + +class TestServiceHelperFunctions: + """Test cases for helper functions in service.py""" + + def test_create_term_with_uri_otype(self): + """Test create_term creates IRI Term for otype='U'""" + from trustgraph.query.triples.cassandra.service import create_term + from trustgraph.schema import IRI + + term = create_term('http://example.org/Alice', otype='U') + + assert term.type == IRI + assert term.iri == 'http://example.org/Alice' + + def test_create_term_with_literal_otype(self): + """Test create_term creates LITERAL Term for otype='L'""" + from trustgraph.query.triples.cassandra.service import create_term + from trustgraph.schema import LITERAL + + term = create_term('Alice Smith', otype='L', dtype='xsd:string', lang='en') + + assert term.type == LITERAL + assert term.value == 'Alice Smith' + assert term.datatype == 'xsd:string' + assert term.language == 'en' + + def test_create_term_with_triple_otype(self): + """Test create_term creates IRI Term for otype='T'""" + from trustgraph.query.triples.cassandra.service import create_term + from trustgraph.schema import IRI + + term = create_term('http://example.org/statement1', otype='T') + + assert term.type == IRI + assert term.iri == 'http://example.org/statement1' + + def test_create_term_heuristic_fallback_uri(self): + """Test create_term uses URL heuristic when otype not provided""" + from trustgraph.query.triples.cassandra.service import create_term + from trustgraph.schema import IRI + + term = create_term('http://example.org/Alice') + + assert term.type == IRI + assert term.iri == 'http://example.org/Alice' + + def test_create_term_heuristic_fallback_literal(self): + """Test create_term uses literal heuristic when otype not provided""" + from trustgraph.query.triples.cassandra.service import create_term + from trustgraph.schema import LITERAL + + term = create_term('Alice Smith') + + assert term.type == LITERAL + assert term.value == 'Alice Smith' diff --git a/tests/unit/test_query/test_triples_cassandra_query.py b/tests/unit/test_query/test_triples_cassandra_query.py index 0c5dc29c..d755faf4 100644 --- a/tests/unit/test_query/test_triples_cassandra_query.py +++ b/tests/unit/test_query/test_triples_cassandra_query.py @@ -70,25 +70,30 @@ class TestCassandraQueryProcessor: assert result.type == LITERAL @pytest.mark.asyncio - @patch('trustgraph.query.triples.cassandra.service.KnowledgeGraph') - async def test_query_triples_spo_query(self, mock_trustgraph): + @patch('trustgraph.query.triples.cassandra.service.get_knowledge_graph_class') + async def test_query_triples_spo_query(self, mock_get_kg_class): """Test querying triples with subject, predicate, and object specified""" from trustgraph.schema import TriplesQueryRequest, Term, IRI, LITERAL - - # Setup mock TrustGraph + + # Setup mock TrustGraph via factory function mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class # SPO query returns a list of results (with mock graph attribute) mock_result = MagicMock() - mock_result.g = None + mock_result.g = '' + mock_result.otype = None + mock_result.dtype = None + mock_result.lang = None + mock_result.o = 'test_object' mock_tg_instance.get_spo.return_value = [mock_result] - + processor = Processor( taskgroup=MagicMock(), id='test-cassandra-query', cassandra_host='localhost' ) - + # Create query request with all SPO values query = TriplesQueryRequest( user='test_user', @@ -98,20 +103,20 @@ class TestCassandraQueryProcessor: o=Term(type=LITERAL, value='test_object'), limit=100 ) - + result = await processor.query_triples(query) - + # Verify KnowledgeGraph was created with correct parameters - mock_trustgraph.assert_called_once_with( + mock_kg_class.assert_called_once_with( hosts=['localhost'], keyspace='test_user' ) - + # Verify get_spo was called with correct parameters mock_tg_instance.get_spo.assert_called_once_with( 'test_collection', 'test_subject', 'test_predicate', 'test_object', g=None, limit=100 ) - + # Verify result contains the queried triple assert len(result) == 1 assert result[0].s.value == 'test_subject' @@ -146,21 +151,26 @@ class TestCassandraQueryProcessor: assert processor.table is None @pytest.mark.asyncio - @patch('trustgraph.query.triples.cassandra.service.KnowledgeGraph') - async def test_query_triples_sp_pattern(self, mock_trustgraph): + @patch('trustgraph.query.triples.cassandra.service.get_knowledge_graph_class') + async def test_query_triples_sp_pattern(self, mock_get_kg_class): """Test SP query pattern (subject and predicate, no object)""" from trustgraph.schema import TriplesQueryRequest, Term, IRI, LITERAL - - # Setup mock TrustGraph and response + + # Setup mock TrustGraph via factory function mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance - + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class + mock_result = MagicMock() mock_result.o = 'result_object' + mock_result.g = '' + mock_result.otype = None + mock_result.dtype = None + mock_result.lang = None mock_tg_instance.get_sp.return_value = [mock_result] - + processor = Processor(taskgroup=MagicMock()) - + query = TriplesQueryRequest( user='test_user', collection='test_collection', @@ -169,9 +179,9 @@ class TestCassandraQueryProcessor: o=None, limit=50 ) - + result = await processor.query_triples(query) - + mock_tg_instance.get_sp.assert_called_once_with('test_collection', 'test_subject', 'test_predicate', g=None, limit=50) assert len(result) == 1 assert result[0].s.value == 'test_subject' @@ -179,21 +189,26 @@ class TestCassandraQueryProcessor: assert result[0].o.value == 'result_object' @pytest.mark.asyncio - @patch('trustgraph.query.triples.cassandra.service.KnowledgeGraph') - async def test_query_triples_s_pattern(self, mock_trustgraph): + @patch('trustgraph.query.triples.cassandra.service.get_knowledge_graph_class') + async def test_query_triples_s_pattern(self, mock_get_kg_class): """Test S query pattern (subject only)""" from trustgraph.schema import TriplesQueryRequest, Term, IRI, LITERAL - + mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance - + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class + mock_result = MagicMock() mock_result.p = 'result_predicate' mock_result.o = 'result_object' + mock_result.g = '' + mock_result.otype = None + mock_result.dtype = None + mock_result.lang = None mock_tg_instance.get_s.return_value = [mock_result] - + processor = Processor(taskgroup=MagicMock()) - + query = TriplesQueryRequest( user='test_user', collection='test_collection', @@ -202,9 +217,9 @@ class TestCassandraQueryProcessor: o=None, limit=25 ) - + result = await processor.query_triples(query) - + mock_tg_instance.get_s.assert_called_once_with('test_collection', 'test_subject', g=None, limit=25) assert len(result) == 1 assert result[0].s.value == 'test_subject' @@ -212,21 +227,26 @@ class TestCassandraQueryProcessor: assert result[0].o.value == 'result_object' @pytest.mark.asyncio - @patch('trustgraph.query.triples.cassandra.service.KnowledgeGraph') - async def test_query_triples_p_pattern(self, mock_trustgraph): + @patch('trustgraph.query.triples.cassandra.service.get_knowledge_graph_class') + async def test_query_triples_p_pattern(self, mock_get_kg_class): """Test P query pattern (predicate only)""" from trustgraph.schema import TriplesQueryRequest, Term, IRI, LITERAL - + mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance - + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class + mock_result = MagicMock() mock_result.s = 'result_subject' mock_result.o = 'result_object' + mock_result.g = '' + mock_result.otype = None + mock_result.dtype = None + mock_result.lang = None mock_tg_instance.get_p.return_value = [mock_result] - + processor = Processor(taskgroup=MagicMock()) - + query = TriplesQueryRequest( user='test_user', collection='test_collection', @@ -235,9 +255,9 @@ class TestCassandraQueryProcessor: o=None, limit=10 ) - + result = await processor.query_triples(query) - + mock_tg_instance.get_p.assert_called_once_with('test_collection', 'test_predicate', g=None, limit=10) assert len(result) == 1 assert result[0].s.value == 'result_subject' @@ -245,21 +265,26 @@ class TestCassandraQueryProcessor: assert result[0].o.value == 'result_object' @pytest.mark.asyncio - @patch('trustgraph.query.triples.cassandra.service.KnowledgeGraph') - async def test_query_triples_o_pattern(self, mock_trustgraph): + @patch('trustgraph.query.triples.cassandra.service.get_knowledge_graph_class') + async def test_query_triples_o_pattern(self, mock_get_kg_class): """Test O query pattern (object only)""" from trustgraph.schema import TriplesQueryRequest, Term, IRI, LITERAL - + mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance - + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class + mock_result = MagicMock() mock_result.s = 'result_subject' mock_result.p = 'result_predicate' + mock_result.g = '' + mock_result.otype = None + mock_result.dtype = None + mock_result.lang = None mock_tg_instance.get_o.return_value = [mock_result] - + processor = Processor(taskgroup=MagicMock()) - + query = TriplesQueryRequest( user='test_user', collection='test_collection', @@ -268,9 +293,9 @@ class TestCassandraQueryProcessor: o=Term(type=LITERAL, value='test_object'), limit=75 ) - + result = await processor.query_triples(query) - + mock_tg_instance.get_o.assert_called_once_with('test_collection', 'test_object', g=None, limit=75) assert len(result) == 1 assert result[0].s.value == 'result_subject' @@ -278,22 +303,27 @@ class TestCassandraQueryProcessor: assert result[0].o.value == 'test_object' @pytest.mark.asyncio - @patch('trustgraph.query.triples.cassandra.service.KnowledgeGraph') - async def test_query_triples_get_all_pattern(self, mock_trustgraph): + @patch('trustgraph.query.triples.cassandra.service.get_knowledge_graph_class') + async def test_query_triples_get_all_pattern(self, mock_get_kg_class): """Test query pattern with no constraints (get all)""" from trustgraph.schema import TriplesQueryRequest - + mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance - + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class + mock_result = MagicMock() mock_result.s = 'all_subject' mock_result.p = 'all_predicate' mock_result.o = 'all_object' + mock_result.g = '' + mock_result.otype = None + mock_result.dtype = None + mock_result.lang = None mock_tg_instance.get_all.return_value = [mock_result] - + processor = Processor(taskgroup=MagicMock()) - + query = TriplesQueryRequest( user='test_user', collection='test_collection', @@ -302,9 +332,9 @@ class TestCassandraQueryProcessor: o=None, limit=1000 ) - + result = await processor.query_triples(query) - + mock_tg_instance.get_all.assert_called_once_with('test_collection', limit=1000) assert len(result) == 1 assert result[0].s.value == 'all_subject' @@ -378,16 +408,21 @@ class TestCassandraQueryProcessor: mock_launch.assert_called_once_with(default_ident, '\nTriples query service. Input is a (s, p, o, g) quad pattern, some values may be\nnull. Output is a list of quads.\n') @pytest.mark.asyncio - @patch('trustgraph.query.triples.cassandra.service.KnowledgeGraph') - async def test_query_triples_with_authentication(self, mock_trustgraph): + @patch('trustgraph.query.triples.cassandra.service.get_knowledge_graph_class') + async def test_query_triples_with_authentication(self, mock_get_kg_class): """Test querying with username and password authentication""" from trustgraph.schema import TriplesQueryRequest, Term, IRI, LITERAL - + mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class # SPO query returns a list of results mock_result = MagicMock() - mock_result.g = None + mock_result.g = '' + mock_result.otype = None + mock_result.dtype = None + mock_result.lang = None + mock_result.o = 'test_object' mock_tg_instance.get_spo.return_value = [mock_result] processor = Processor( @@ -395,7 +430,7 @@ class TestCassandraQueryProcessor: cassandra_username='authuser', cassandra_password='authpass' ) - + query = TriplesQueryRequest( user='test_user', collection='test_collection', @@ -404,11 +439,11 @@ class TestCassandraQueryProcessor: o=Term(type=LITERAL, value='test_object'), limit=100 ) - + await processor.query_triples(query) - + # Verify KnowledgeGraph was created with authentication - mock_trustgraph.assert_called_once_with( + mock_kg_class.assert_called_once_with( hosts=['cassandra'], # Updated default keyspace='test_user', username='authuser', @@ -416,16 +451,21 @@ class TestCassandraQueryProcessor: ) @pytest.mark.asyncio - @patch('trustgraph.query.triples.cassandra.service.KnowledgeGraph') - async def test_query_triples_table_reuse(self, mock_trustgraph): + @patch('trustgraph.query.triples.cassandra.service.get_knowledge_graph_class') + async def test_query_triples_table_reuse(self, mock_get_kg_class): """Test that TrustGraph is reused for same table""" from trustgraph.schema import TriplesQueryRequest, Term, IRI, LITERAL - + mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class # SPO query returns a list of results mock_result = MagicMock() - mock_result.g = None + mock_result.g = '' + mock_result.otype = None + mock_result.dtype = None + mock_result.lang = None + mock_result.o = 'test_object' mock_tg_instance.get_spo.return_value = [mock_result] processor = Processor(taskgroup=MagicMock()) @@ -441,24 +481,36 @@ class TestCassandraQueryProcessor: # First query should create TrustGraph await processor.query_triples(query) - assert mock_trustgraph.call_count == 1 + assert mock_kg_class.call_count == 1 # Second query with same table should reuse TrustGraph await processor.query_triples(query) - assert mock_trustgraph.call_count == 1 # Should not increase + assert mock_kg_class.call_count == 1 # Should not increase @pytest.mark.asyncio - @patch('trustgraph.query.triples.cassandra.service.KnowledgeGraph') - async def test_query_triples_table_switching(self, mock_trustgraph): + @patch('trustgraph.query.triples.cassandra.service.get_knowledge_graph_class') + async def test_query_triples_table_switching(self, mock_get_kg_class): """Test table switching creates new TrustGraph""" from trustgraph.schema import TriplesQueryRequest, Term, IRI, LITERAL - + mock_tg_instance1 = MagicMock() mock_tg_instance2 = MagicMock() - mock_trustgraph.side_effect = [mock_tg_instance1, mock_tg_instance2] - + mock_kg_class = MagicMock(side_effect=[mock_tg_instance1, mock_tg_instance2]) + mock_get_kg_class.return_value = mock_kg_class + + # Setup mock results for both instances + mock_result = MagicMock() + mock_result.g = '' + mock_result.otype = None + mock_result.dtype = None + mock_result.lang = None + mock_result.p = 'p' + mock_result.o = 'o' + mock_tg_instance1.get_s.return_value = [mock_result] + mock_tg_instance2.get_s.return_value = [mock_result] + processor = Processor(taskgroup=MagicMock()) - + # First query query1 = TriplesQueryRequest( user='user1', @@ -468,10 +520,10 @@ class TestCassandraQueryProcessor: o=None, limit=100 ) - + await processor.query_triples(query1) assert processor.table == 'user1' - + # Second query with different table query2 = TriplesQueryRequest( user='user2', @@ -481,25 +533,26 @@ class TestCassandraQueryProcessor: o=None, limit=100 ) - + await processor.query_triples(query2) assert processor.table == 'user2' - + # Verify TrustGraph was created twice - assert mock_trustgraph.call_count == 2 + assert mock_kg_class.call_count == 2 @pytest.mark.asyncio - @patch('trustgraph.query.triples.cassandra.service.KnowledgeGraph') - async def test_query_triples_exception_handling(self, mock_trustgraph): + @patch('trustgraph.query.triples.cassandra.service.get_knowledge_graph_class') + async def test_query_triples_exception_handling(self, mock_get_kg_class): """Test exception handling during query execution""" from trustgraph.schema import TriplesQueryRequest, Term, IRI, LITERAL - + mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class mock_tg_instance.get_spo.side_effect = Exception("Query failed") - + processor = Processor(taskgroup=MagicMock()) - + query = TriplesQueryRequest( user='test_user', collection='test_collection', @@ -508,28 +561,37 @@ class TestCassandraQueryProcessor: o=Term(type=LITERAL, value='test_object'), limit=100 ) - + with pytest.raises(Exception, match="Query failed"): await processor.query_triples(query) @pytest.mark.asyncio - @patch('trustgraph.query.triples.cassandra.service.KnowledgeGraph') - async def test_query_triples_multiple_results(self, mock_trustgraph): + @patch('trustgraph.query.triples.cassandra.service.get_knowledge_graph_class') + async def test_query_triples_multiple_results(self, mock_get_kg_class): """Test query returning multiple results""" from trustgraph.schema import TriplesQueryRequest, Term, IRI, LITERAL - + mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance - + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class + # Mock multiple results mock_result1 = MagicMock() mock_result1.o = 'object1' + mock_result1.g = '' + mock_result1.otype = None + mock_result1.dtype = None + mock_result1.lang = None mock_result2 = MagicMock() mock_result2.o = 'object2' + mock_result2.g = '' + mock_result2.otype = None + mock_result2.dtype = None + mock_result2.lang = None mock_tg_instance.get_sp.return_value = [mock_result1, mock_result2] - + processor = Processor(taskgroup=MagicMock()) - + query = TriplesQueryRequest( user='test_user', collection='test_collection', @@ -538,9 +600,9 @@ class TestCassandraQueryProcessor: o=None, limit=100 ) - + result = await processor.query_triples(query) - + assert len(result) == 2 assert result[0].o.value == 'object1' assert result[1].o.value == 'object2' @@ -550,16 +612,21 @@ class TestCassandraQueryPerformanceOptimizations: """Test cases for multi-table performance optimizations in query service""" @pytest.mark.asyncio - @patch('trustgraph.query.triples.cassandra.service.KnowledgeGraph') - async def test_get_po_query_optimization(self, mock_trustgraph): + @patch('trustgraph.query.triples.cassandra.service.get_knowledge_graph_class') + async def test_get_po_query_optimization(self, mock_get_kg_class): """Test that get_po queries use optimized table (no ALLOW FILTERING)""" from trustgraph.schema import TriplesQueryRequest, Term, IRI, LITERAL mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class mock_result = MagicMock() mock_result.s = 'result_subject' + mock_result.g = '' + mock_result.otype = None + mock_result.dtype = None + mock_result.lang = None mock_tg_instance.get_po.return_value = [mock_result] processor = Processor(taskgroup=MagicMock()) @@ -587,16 +654,21 @@ class TestCassandraQueryPerformanceOptimizations: assert result[0].o.value == 'test_object' @pytest.mark.asyncio - @patch('trustgraph.query.triples.cassandra.service.KnowledgeGraph') - async def test_get_os_query_optimization(self, mock_trustgraph): + @patch('trustgraph.query.triples.cassandra.service.get_knowledge_graph_class') + async def test_get_os_query_optimization(self, mock_get_kg_class): """Test that get_os queries use optimized table (no ALLOW FILTERING)""" from trustgraph.schema import TriplesQueryRequest, Term, IRI, LITERAL mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class mock_result = MagicMock() mock_result.p = 'result_predicate' + mock_result.g = '' + mock_result.otype = None + mock_result.dtype = None + mock_result.lang = None mock_tg_instance.get_os.return_value = [mock_result] processor = Processor(taskgroup=MagicMock()) @@ -624,13 +696,14 @@ class TestCassandraQueryPerformanceOptimizations: assert result[0].o.value == 'test_object' @pytest.mark.asyncio - @patch('trustgraph.query.triples.cassandra.service.KnowledgeGraph') - async def test_all_query_patterns_use_correct_tables(self, mock_trustgraph): + @patch('trustgraph.query.triples.cassandra.service.get_knowledge_graph_class') + async def test_all_query_patterns_use_correct_tables(self, mock_get_kg_class): """Test that all query patterns route to their optimal tables""" from trustgraph.schema import TriplesQueryRequest, Term, IRI, LITERAL mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class # Mock empty results for all queries mock_tg_instance.get_all.return_value = [] @@ -696,19 +769,24 @@ class TestCassandraQueryPerformanceOptimizations: # Mode is determined in KnowledgeGraph initialization @pytest.mark.asyncio - @patch('trustgraph.query.triples.cassandra.service.KnowledgeGraph') - async def test_performance_critical_po_query_no_filtering(self, mock_trustgraph): + @patch('trustgraph.query.triples.cassandra.service.get_knowledge_graph_class') + async def test_performance_critical_po_query_no_filtering(self, mock_get_kg_class): """Test the performance-critical PO query that eliminates ALLOW FILTERING""" from trustgraph.schema import TriplesQueryRequest, Term, IRI, LITERAL mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class # Mock multiple subjects for the same predicate-object pair mock_results = [] for i in range(5): mock_result = MagicMock() mock_result.s = f'subject_{i}' + mock_result.g = '' + mock_result.otype = None + mock_result.dtype = None + mock_result.lang = None mock_results.append(mock_result) mock_tg_instance.get_po.return_value = mock_results diff --git a/tests/unit/test_storage/test_triples_cassandra_storage.py b/tests/unit/test_storage/test_triples_cassandra_storage.py index 3fdff6b9..aa7a3383 100644 --- a/tests/unit/test_storage/test_triples_cassandra_storage.py +++ b/tests/unit/test_storage/test_triples_cassandra_storage.py @@ -6,7 +6,7 @@ import pytest from unittest.mock import MagicMock, patch, AsyncMock from trustgraph.storage.triples.cassandra.write import Processor -from trustgraph.schema import Triple, LITERAL +from trustgraph.schema import Triple, LITERAL, IRI from trustgraph.direct.cassandra_kg import DEFAULT_GRAPH @@ -87,29 +87,30 @@ class TestCassandraStorageProcessor: assert processor.cassandra_username == 'new-user' # Only cassandra_* params work @pytest.mark.asyncio - @patch('trustgraph.storage.triples.cassandra.write.KnowledgeGraph') - async def test_table_switching_with_auth(self, mock_trustgraph): + @patch('trustgraph.storage.triples.cassandra.write.get_knowledge_graph_class') + async def test_table_switching_with_auth(self, mock_get_kg_class): """Test table switching logic when authentication is provided""" taskgroup_mock = MagicMock() mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance - + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class + processor = Processor( taskgroup=taskgroup_mock, cassandra_username='testuser', cassandra_password='testpass' ) - + # Create mock message mock_message = MagicMock() mock_message.metadata.user = 'user1' mock_message.metadata.collection = 'collection1' mock_message.triples = [] - + await processor.store_triples(mock_message) - + # Verify KnowledgeGraph was called with auth parameters - mock_trustgraph.assert_called_once_with( + mock_kg_class.assert_called_once_with( hosts=['cassandra'], # Updated default keyspace='user1', username='testuser', @@ -118,81 +119,92 @@ class TestCassandraStorageProcessor: assert processor.table == 'user1' @pytest.mark.asyncio - @patch('trustgraph.storage.triples.cassandra.write.KnowledgeGraph') - async def test_table_switching_without_auth(self, mock_trustgraph): + @patch('trustgraph.storage.triples.cassandra.write.get_knowledge_graph_class') + async def test_table_switching_without_auth(self, mock_get_kg_class): """Test table switching logic when no authentication is provided""" taskgroup_mock = MagicMock() mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance - + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class + processor = Processor(taskgroup=taskgroup_mock) - + # Create mock message mock_message = MagicMock() mock_message.metadata.user = 'user2' mock_message.metadata.collection = 'collection2' mock_message.triples = [] - + await processor.store_triples(mock_message) - + # Verify KnowledgeGraph was called without auth parameters - mock_trustgraph.assert_called_once_with( + mock_kg_class.assert_called_once_with( hosts=['cassandra'], # Updated default keyspace='user2' ) assert processor.table == 'user2' @pytest.mark.asyncio - @patch('trustgraph.storage.triples.cassandra.write.KnowledgeGraph') - async def test_table_reuse_when_same(self, mock_trustgraph): + @patch('trustgraph.storage.triples.cassandra.write.get_knowledge_graph_class') + async def test_table_reuse_when_same(self, mock_get_kg_class): """Test that TrustGraph is not recreated when table hasn't changed""" taskgroup_mock = MagicMock() mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance - + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class + processor = Processor(taskgroup=taskgroup_mock) - + # Create mock message mock_message = MagicMock() mock_message.metadata.user = 'user1' mock_message.metadata.collection = 'collection1' mock_message.triples = [] - + # First call should create TrustGraph await processor.store_triples(mock_message) - assert mock_trustgraph.call_count == 1 - + assert mock_kg_class.call_count == 1 + # Second call with same table should reuse TrustGraph await processor.store_triples(mock_message) - assert mock_trustgraph.call_count == 1 # Should not increase + assert mock_kg_class.call_count == 1 # Should not increase @pytest.mark.asyncio - @patch('trustgraph.storage.triples.cassandra.write.KnowledgeGraph') - async def test_triple_insertion(self, mock_trustgraph): + @patch('trustgraph.storage.triples.cassandra.write.get_knowledge_graph_class') + async def test_triple_insertion(self, mock_get_kg_class): """Test that triples are properly inserted into Cassandra""" taskgroup_mock = MagicMock() mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance - + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class + processor = Processor(taskgroup=taskgroup_mock) - + # Create mock triples with proper Term structure triple1 = MagicMock() triple1.s.type = LITERAL triple1.s.value = 'subject1' + triple1.s.datatype = '' + triple1.s.language = '' triple1.p.type = LITERAL triple1.p.value = 'predicate1' triple1.o.type = LITERAL triple1.o.value = 'object1' + triple1.o.datatype = '' + triple1.o.language = '' triple1.g = None triple2 = MagicMock() triple2.s.type = LITERAL triple2.s.value = 'subject2' + triple2.s.datatype = '' + triple2.s.language = '' triple2.p.type = LITERAL triple2.p.value = 'predicate2' triple2.o.type = LITERAL triple2.o.value = 'object2' + triple2.o.datatype = '' + triple2.o.language = '' triple2.g = None # Create mock message @@ -203,51 +215,59 @@ class TestCassandraStorageProcessor: await processor.store_triples(mock_message) - # Verify both triples were inserted (with g= parameter) + # Verify both triples were inserted (with g=, otype=, dtype=, lang= parameters) assert mock_tg_instance.insert.call_count == 2 - mock_tg_instance.insert.assert_any_call('collection1', 'subject1', 'predicate1', 'object1', g=DEFAULT_GRAPH) - mock_tg_instance.insert.assert_any_call('collection1', 'subject2', 'predicate2', 'object2', g=DEFAULT_GRAPH) + mock_tg_instance.insert.assert_any_call( + 'collection1', 'subject1', 'predicate1', 'object1', + g=DEFAULT_GRAPH, otype='L', dtype='', lang='' + ) + mock_tg_instance.insert.assert_any_call( + 'collection1', 'subject2', 'predicate2', 'object2', + g=DEFAULT_GRAPH, otype='L', dtype='', lang='' + ) @pytest.mark.asyncio - @patch('trustgraph.storage.triples.cassandra.write.KnowledgeGraph') - async def test_triple_insertion_with_empty_list(self, mock_trustgraph): + @patch('trustgraph.storage.triples.cassandra.write.get_knowledge_graph_class') + async def test_triple_insertion_with_empty_list(self, mock_get_kg_class): """Test behavior when message has no triples""" taskgroup_mock = MagicMock() mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance - + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class + processor = Processor(taskgroup=taskgroup_mock) - + # Create mock message with empty triples mock_message = MagicMock() mock_message.metadata.user = 'user1' mock_message.metadata.collection = 'collection1' mock_message.triples = [] - + await processor.store_triples(mock_message) - + # Verify no triples were inserted mock_tg_instance.insert.assert_not_called() @pytest.mark.asyncio - @patch('trustgraph.storage.triples.cassandra.write.KnowledgeGraph') + @patch('trustgraph.storage.triples.cassandra.write.get_knowledge_graph_class') @patch('trustgraph.storage.triples.cassandra.write.time.sleep') - async def test_exception_handling_with_retry(self, mock_sleep, mock_trustgraph): + async def test_exception_handling_with_retry(self, mock_sleep, mock_get_kg_class): """Test exception handling during TrustGraph creation""" taskgroup_mock = MagicMock() - mock_trustgraph.side_effect = Exception("Connection failed") - + mock_kg_class = MagicMock(side_effect=Exception("Connection failed")) + mock_get_kg_class.return_value = mock_kg_class + processor = Processor(taskgroup=taskgroup_mock) - + # Create mock message mock_message = MagicMock() mock_message.metadata.user = 'user1' mock_message.metadata.collection = 'collection1' mock_message.triples = [] - + with pytest.raises(Exception, match="Connection failed"): await processor.store_triples(mock_message) - + # Verify sleep was called before re-raising mock_sleep.assert_called_once_with(1) @@ -335,57 +355,63 @@ class TestCassandraStorageProcessor: mock_launch.assert_called_once_with(default_ident, '\nGraph writer. Input is graph edge. Writes edges to Cassandra graph.\n') @pytest.mark.asyncio - @patch('trustgraph.storage.triples.cassandra.write.KnowledgeGraph') - async def test_store_triples_table_switching_between_different_tables(self, mock_trustgraph): + @patch('trustgraph.storage.triples.cassandra.write.get_knowledge_graph_class') + async def test_store_triples_table_switching_between_different_tables(self, mock_get_kg_class): """Test table switching when different tables are used in sequence""" taskgroup_mock = MagicMock() mock_tg_instance1 = MagicMock() mock_tg_instance2 = MagicMock() - mock_trustgraph.side_effect = [mock_tg_instance1, mock_tg_instance2] - + mock_kg_class = MagicMock(side_effect=[mock_tg_instance1, mock_tg_instance2]) + mock_get_kg_class.return_value = mock_kg_class + processor = Processor(taskgroup=taskgroup_mock) - + # First message with table1 mock_message1 = MagicMock() mock_message1.metadata.user = 'user1' mock_message1.metadata.collection = 'collection1' mock_message1.triples = [] - + await processor.store_triples(mock_message1) assert processor.table == 'user1' assert processor.tg == mock_tg_instance1 - + # Second message with different table mock_message2 = MagicMock() mock_message2.metadata.user = 'user2' mock_message2.metadata.collection = 'collection2' mock_message2.triples = [] - + await processor.store_triples(mock_message2) assert processor.table == 'user2' assert processor.tg == mock_tg_instance2 - + # Verify TrustGraph was created twice for different tables - assert mock_trustgraph.call_count == 2 + assert mock_kg_class.call_count == 2 @pytest.mark.asyncio - @patch('trustgraph.storage.triples.cassandra.write.KnowledgeGraph') - async def test_store_triples_with_special_characters_in_values(self, mock_trustgraph): + @patch('trustgraph.storage.triples.cassandra.write.get_knowledge_graph_class') + async def test_store_triples_with_special_characters_in_values(self, mock_get_kg_class): """Test storing triples with special characters and unicode""" taskgroup_mock = MagicMock() mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance - + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class + processor = Processor(taskgroup=taskgroup_mock) - + # Create triple with special characters and proper Term structure triple = MagicMock() triple.s.type = LITERAL triple.s.value = 'subject with spaces & symbols' + triple.s.datatype = '' + triple.s.language = '' triple.p.type = LITERAL triple.p.value = 'predicate:with/colons' triple.o.type = LITERAL triple.o.value = 'object with "quotes" and unicode: ñáéíóú' + triple.o.datatype = '' + triple.o.language = '' triple.g = None mock_message = MagicMock() @@ -401,31 +427,35 @@ class TestCassandraStorageProcessor: 'subject with spaces & symbols', 'predicate:with/colons', 'object with "quotes" and unicode: ñáéíóú', - g=DEFAULT_GRAPH + g=DEFAULT_GRAPH, + otype='L', + dtype='', + lang='' ) @pytest.mark.asyncio - @patch('trustgraph.storage.triples.cassandra.write.KnowledgeGraph') - async def test_store_triples_preserves_old_table_on_exception(self, mock_trustgraph): + @patch('trustgraph.storage.triples.cassandra.write.get_knowledge_graph_class') + async def test_store_triples_preserves_old_table_on_exception(self, mock_get_kg_class): """Test that table remains unchanged when TrustGraph creation fails""" taskgroup_mock = MagicMock() - + processor = Processor(taskgroup=taskgroup_mock) - + # Set an initial table processor.table = ('old_user', 'old_collection') - + # Mock TrustGraph to raise exception - mock_trustgraph.side_effect = Exception("Connection failed") - + mock_kg_class = MagicMock(side_effect=Exception("Connection failed")) + mock_get_kg_class.return_value = mock_kg_class + mock_message = MagicMock() mock_message.metadata.user = 'new_user' mock_message.metadata.collection = 'new_collection' mock_message.triples = [] - + with pytest.raises(Exception, match="Connection failed"): await processor.store_triples(mock_message) - + # Table should remain unchanged since self.table = table happens after try/except assert processor.table == ('old_user', 'old_collection') # TrustGraph should be set to None though @@ -436,12 +466,13 @@ class TestCassandraPerformanceOptimizations: """Test cases for multi-table performance optimizations""" @pytest.mark.asyncio - @patch('trustgraph.storage.triples.cassandra.write.KnowledgeGraph') - async def test_legacy_mode_uses_single_table(self, mock_trustgraph): + @patch('trustgraph.storage.triples.cassandra.write.get_knowledge_graph_class') + async def test_legacy_mode_uses_single_table(self, mock_get_kg_class): """Test that legacy mode still works with single table""" taskgroup_mock = MagicMock() mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class with patch.dict('os.environ', {'CASSANDRA_USE_LEGACY': 'true'}): processor = Processor(taskgroup=taskgroup_mock) @@ -454,16 +485,16 @@ class TestCassandraPerformanceOptimizations: await processor.store_triples(mock_message) # Verify KnowledgeGraph instance uses legacy mode - kg_instance = mock_trustgraph.return_value - assert kg_instance is not None + assert mock_tg_instance is not None @pytest.mark.asyncio - @patch('trustgraph.storage.triples.cassandra.write.KnowledgeGraph') - async def test_optimized_mode_uses_multi_table(self, mock_trustgraph): + @patch('trustgraph.storage.triples.cassandra.write.get_knowledge_graph_class') + async def test_optimized_mode_uses_multi_table(self, mock_get_kg_class): """Test that optimized mode uses multi-table schema""" taskgroup_mock = MagicMock() mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class with patch.dict('os.environ', {'CASSANDRA_USE_LEGACY': 'false'}): processor = Processor(taskgroup=taskgroup_mock) @@ -476,16 +507,16 @@ class TestCassandraPerformanceOptimizations: await processor.store_triples(mock_message) # Verify KnowledgeGraph instance is in optimized mode - kg_instance = mock_trustgraph.return_value - assert kg_instance is not None + assert mock_tg_instance is not None @pytest.mark.asyncio - @patch('trustgraph.storage.triples.cassandra.write.KnowledgeGraph') - async def test_batch_write_consistency(self, mock_trustgraph): + @patch('trustgraph.storage.triples.cassandra.write.get_knowledge_graph_class') + async def test_batch_write_consistency(self, mock_get_kg_class): """Test that all tables stay consistent during batch writes""" taskgroup_mock = MagicMock() mock_tg_instance = MagicMock() - mock_trustgraph.return_value = mock_tg_instance + mock_kg_class = MagicMock(return_value=mock_tg_instance) + mock_get_kg_class.return_value = mock_kg_class processor = Processor(taskgroup=taskgroup_mock) @@ -493,10 +524,14 @@ class TestCassandraPerformanceOptimizations: triple = MagicMock() triple.s.type = LITERAL triple.s.value = 'test_subject' + triple.s.datatype = '' + triple.s.language = '' triple.p.type = LITERAL triple.p.value = 'test_predicate' triple.o.type = LITERAL triple.o.value = 'test_object' + triple.o.datatype = '' + triple.o.language = '' triple.g = None mock_message = MagicMock() @@ -509,7 +544,7 @@ class TestCassandraPerformanceOptimizations: # Verify insert was called for the triple (implementation details tested in KnowledgeGraph) mock_tg_instance.insert.assert_called_once_with( 'collection1', 'test_subject', 'test_predicate', 'test_object', - g=DEFAULT_GRAPH + g=DEFAULT_GRAPH, otype='L', dtype='', lang='' ) def test_environment_variable_controls_mode(self): diff --git a/trustgraph-flow/trustgraph/direct/cassandra_kg.py b/trustgraph-flow/trustgraph/direct/cassandra_kg.py index f8a20041..20bd564a 100644 --- a/trustgraph-flow/trustgraph/direct/cassandra_kg.py +++ b/trustgraph-flow/trustgraph/direct/cassandra_kg.py @@ -516,3 +516,590 @@ class KnowledgeGraph: self.cluster.shutdown() if self.cluster in _active_clusters: _active_clusters.remove(self.cluster) + + +class EntityCentricKnowledgeGraph: + """ + Entity-centric Cassandra-backed knowledge graph supporting quads (s, p, o, g). + + Uses 2 tables instead of 7: + - quads_by_entity: every entity knows every quad it participates in + - quads_by_collection: manifest for collection-level queries and deletion + + Supports all 16 query patterns with single-partition reads. + """ + + def __init__( + self, hosts=None, + keyspace="trustgraph", username=None, password=None + ): + + if hosts is None: + hosts = ["localhost"] + + self.keyspace = keyspace + self.username = username + + # 2-table entity-centric schema + self.entity_table = "quads_by_entity" + self.collection_table = "quads_by_collection" + + # Collection metadata tracking + self.collection_metadata_table = "collection_metadata" + + if username and password: + ssl_context = SSLContext(PROTOCOL_TLSv1_2) + auth_provider = PlainTextAuthProvider(username=username, password=password) + self.cluster = Cluster(hosts, auth_provider=auth_provider, ssl_context=ssl_context) + else: + self.cluster = Cluster(hosts) + self.session = self.cluster.connect() + + # Track this cluster globally + _active_clusters.append(self.cluster) + + self.init() + self.prepare_statements() + + def clear(self): + self.session.execute(f""" + drop keyspace if exists {self.keyspace}; + """) + self.init() + + def init(self): + self.session.execute(f""" + create keyspace if not exists {self.keyspace} + with replication = {{ + 'class' : 'SimpleStrategy', + 'replication_factor' : 1 + }}; + """) + + self.session.set_keyspace(self.keyspace) + self.init_entity_centric_schema() + + def init_entity_centric_schema(self): + """Initialize 2-table entity-centric schema""" + + # quads_by_entity: primary data table + # Every entity has a partition containing all quads it participates in + self.session.execute(f""" + CREATE TABLE IF NOT EXISTS {self.entity_table} ( + collection text, + entity text, + role text, + p text, + otype text, + s text, + o text, + d text, + dtype text, + lang text, + PRIMARY KEY ((collection, entity), role, p, otype, s, o, d) + ); + """) + + # quads_by_collection: manifest for collection-level queries and deletion + self.session.execute(f""" + CREATE TABLE IF NOT EXISTS {self.collection_table} ( + collection text, + d text, + s text, + p text, + o text, + otype text, + dtype text, + lang text, + PRIMARY KEY (collection, d, s, p, o) + ); + """) + + # Collection metadata tracking + self.session.execute(f""" + CREATE TABLE IF NOT EXISTS {self.collection_metadata_table} ( + collection text, + created_at timestamp, + PRIMARY KEY (collection) + ); + """) + + logger.info("Entity-centric schema initialized (2 tables + metadata)") + + def prepare_statements(self): + """Prepare statements for entity-centric schema""" + + # Insert statement for quads_by_entity + self.insert_entity_stmt = self.session.prepare( + f"INSERT INTO {self.entity_table} " + "(collection, entity, role, p, otype, s, o, d, dtype, lang) " + "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" + ) + + # Insert statement for quads_by_collection + self.insert_collection_stmt = self.session.prepare( + f"INSERT INTO {self.collection_table} " + "(collection, d, s, p, o, otype, dtype, lang) " + "VALUES (?, ?, ?, ?, ?, ?, ?, ?)" + ) + + # Query statements for quads_by_entity + + # Get all quads for an entity (any role) + self.get_entity_all_stmt = self.session.prepare( + f"SELECT role, p, otype, s, o, d, dtype, lang FROM {self.entity_table} " + "WHERE collection = ? AND entity = ? LIMIT ?" + ) + + # Get quads where entity is subject (role='S') + self.get_entity_as_s_stmt = self.session.prepare( + f"SELECT p, otype, s, o, d, dtype, lang FROM {self.entity_table} " + "WHERE collection = ? AND entity = ? AND role = 'S' LIMIT ?" + ) + + # Get quads where entity is subject with specific predicate + self.get_entity_as_s_p_stmt = self.session.prepare( + f"SELECT otype, s, o, d, dtype, lang FROM {self.entity_table} " + "WHERE collection = ? AND entity = ? AND role = 'S' AND p = ? LIMIT ?" + ) + + # Get quads where entity is subject with specific predicate and otype + self.get_entity_as_s_p_otype_stmt = self.session.prepare( + f"SELECT s, o, d, dtype, lang FROM {self.entity_table} " + "WHERE collection = ? AND entity = ? AND role = 'S' AND p = ? AND otype = ? LIMIT ?" + ) + + # Get quads where entity is predicate (role='P') + self.get_entity_as_p_stmt = self.session.prepare( + f"SELECT p, otype, s, o, d, dtype, lang FROM {self.entity_table} " + "WHERE collection = ? AND entity = ? AND role = 'P' LIMIT ?" + ) + + # Get quads where entity is object (role='O') + self.get_entity_as_o_stmt = self.session.prepare( + f"SELECT p, otype, s, o, d, dtype, lang FROM {self.entity_table} " + "WHERE collection = ? AND entity = ? AND role = 'O' LIMIT ?" + ) + + # Get quads where entity is object with specific predicate + self.get_entity_as_o_p_stmt = self.session.prepare( + f"SELECT otype, s, o, d, dtype, lang FROM {self.entity_table} " + "WHERE collection = ? AND entity = ? AND role = 'O' AND p = ? LIMIT ?" + ) + + # Get quads where entity is graph (role='G') + self.get_entity_as_g_stmt = self.session.prepare( + f"SELECT p, otype, s, o, d, dtype, lang FROM {self.entity_table} " + "WHERE collection = ? AND entity = ? AND role = 'G' LIMIT ?" + ) + + # Query statements for quads_by_collection + + # Get all quads in collection + self.get_collection_all_stmt = self.session.prepare( + f"SELECT d, s, p, o, otype, dtype, lang FROM {self.collection_table} " + "WHERE collection = ? LIMIT ?" + ) + + # Get all quads in a specific graph + self.get_collection_by_graph_stmt = self.session.prepare( + f"SELECT s, p, o, otype, dtype, lang FROM {self.collection_table} " + "WHERE collection = ? AND d = ? LIMIT ?" + ) + + # Delete statements + self.delete_entity_partition_stmt = self.session.prepare( + f"DELETE FROM {self.entity_table} WHERE collection = ? AND entity = ?" + ) + + self.delete_collection_row_stmt = self.session.prepare( + f"DELETE FROM {self.collection_table} WHERE collection = ? AND d = ? AND s = ? AND p = ? AND o = ?" + ) + + logger.info("Prepared statements initialized for entity-centric schema") + + def insert(self, collection, s, p, o, g=None, otype=None, dtype="", lang=""): + """ + Insert a quad into entity-centric tables. + + Writes 4 rows to quads_by_entity (one for each entity role) + 1 row to + quads_by_collection. For literals, only 3 entity rows are written since + literals are not independently queryable entities. + + Args: + collection: Collection/tenant scope + s: Subject (string value) + p: Predicate (string value) + o: Object (string value) + g: Graph/dataset (None for default graph) + otype: Object type - 'U' (URI), 'L' (literal), 'T' (triple) + Auto-detected from o value if not provided + dtype: XSD datatype (for literals) + lang: Language tag (for literals) + """ + # Default graph stored as empty string + if g is None: + g = DEFAULT_GRAPH + + # Auto-detect otype if not provided (backwards compatibility) + if otype is None: + if o.startswith("http://") or o.startswith("https://"): + otype = "U" + else: + otype = "L" + + batch = BatchStatement() + + # Write row for subject entity (role='S') + batch.add(self.insert_entity_stmt, ( + collection, s, 'S', p, otype, s, o, g, dtype, lang + )) + + # Write row for predicate entity (role='P') + batch.add(self.insert_entity_stmt, ( + collection, p, 'P', p, otype, s, o, g, dtype, lang + )) + + # Write row for object entity (role='O') - only for URIs, not literals + if otype == 'U' or otype == 'T': + batch.add(self.insert_entity_stmt, ( + collection, o, 'O', p, otype, s, o, g, dtype, lang + )) + + # Write row for graph entity (role='G') - only for non-default graphs + if g != DEFAULT_GRAPH: + batch.add(self.insert_entity_stmt, ( + collection, g, 'G', p, otype, s, o, g, dtype, lang + )) + + # Write row to quads_by_collection + batch.add(self.insert_collection_stmt, ( + collection, g, s, p, o, otype, dtype, lang + )) + + self.session.execute(batch) + + # ======================================================================== + # Query methods + # g=None means default graph, g="*" means all graphs + # Results include otype, dtype, lang for proper Term reconstruction + # ======================================================================== + + def get_all(self, collection, limit=50): + """Get all quads in collection""" + return self.session.execute(self.get_collection_all_stmt, (collection, limit)) + + def get_s(self, collection, s, g=None, limit=10): + """ + Query by subject. Returns quads where s is the subject. + g=None: default graph, g='*': all graphs + """ + rows = self.session.execute(self.get_entity_as_s_stmt, (collection, s, limit)) + + results = [] + for row in rows: + d = row.d if hasattr(row, 'd') else DEFAULT_GRAPH + # Filter by graph if specified + if g is None or g == DEFAULT_GRAPH: + if d != DEFAULT_GRAPH: + continue + elif g != GRAPH_WILDCARD and d != g: + continue + + results.append(QuadResult( + s=row.s, p=row.p, o=row.o, g=d, + otype=row.otype, dtype=row.dtype, lang=row.lang + )) + + return results + + def get_p(self, collection, p, g=None, limit=10): + """Query by predicate""" + rows = self.session.execute(self.get_entity_as_p_stmt, (collection, p, limit)) + + results = [] + for row in rows: + d = row.d if hasattr(row, 'd') else DEFAULT_GRAPH + if g is None or g == DEFAULT_GRAPH: + if d != DEFAULT_GRAPH: + continue + elif g != GRAPH_WILDCARD and d != g: + continue + + results.append(QuadResult( + s=row.s, p=row.p, o=row.o, g=d, + otype=row.otype, dtype=row.dtype, lang=row.lang + )) + + return results + + def get_o(self, collection, o, g=None, limit=10): + """Query by object""" + rows = self.session.execute(self.get_entity_as_o_stmt, (collection, o, limit)) + + results = [] + for row in rows: + d = row.d if hasattr(row, 'd') else DEFAULT_GRAPH + if g is None or g == DEFAULT_GRAPH: + if d != DEFAULT_GRAPH: + continue + elif g != GRAPH_WILDCARD and d != g: + continue + + results.append(QuadResult( + s=row.s, p=row.p, o=row.o, g=d, + otype=row.otype, dtype=row.dtype, lang=row.lang + )) + + return results + + def get_sp(self, collection, s, p, g=None, limit=10): + """Query by subject and predicate""" + rows = self.session.execute(self.get_entity_as_s_p_stmt, (collection, s, p, limit)) + + results = [] + for row in rows: + d = row.d if hasattr(row, 'd') else DEFAULT_GRAPH + if g is None or g == DEFAULT_GRAPH: + if d != DEFAULT_GRAPH: + continue + elif g != GRAPH_WILDCARD and d != g: + continue + + results.append(QuadResult( + s=s, p=p, o=row.o, g=d, + otype=row.otype, dtype=row.dtype, lang=row.lang + )) + + return results + + def get_po(self, collection, p, o, g=None, limit=10): + """Query by predicate and object""" + rows = self.session.execute(self.get_entity_as_o_p_stmt, (collection, o, p, limit)) + + results = [] + for row in rows: + d = row.d if hasattr(row, 'd') else DEFAULT_GRAPH + if g is None or g == DEFAULT_GRAPH: + if d != DEFAULT_GRAPH: + continue + elif g != GRAPH_WILDCARD and d != g: + continue + + results.append(QuadResult( + s=row.s, p=p, o=o, g=d, + otype=row.otype, dtype=row.dtype, lang=row.lang + )) + + return results + + def get_os(self, collection, o, s, g=None, limit=10): + """Query by object and subject""" + # Use subject partition with role='S', filter by o + rows = self.session.execute(self.get_entity_as_s_stmt, (collection, s, limit)) + + results = [] + for row in rows: + if row.o != o: + continue + + d = row.d if hasattr(row, 'd') else DEFAULT_GRAPH + if g is None or g == DEFAULT_GRAPH: + if d != DEFAULT_GRAPH: + continue + elif g != GRAPH_WILDCARD and d != g: + continue + + results.append(QuadResult( + s=s, p=row.p, o=o, g=d, + otype=row.otype, dtype=row.dtype, lang=row.lang + )) + + return results + + def get_spo(self, collection, s, p, o, g=None, limit=10): + """Query by subject, predicate, object (find which graphs)""" + rows = self.session.execute(self.get_entity_as_s_p_stmt, (collection, s, p, limit)) + + results = [] + for row in rows: + if row.o != o: + continue + + d = row.d if hasattr(row, 'd') else DEFAULT_GRAPH + if g is None or g == DEFAULT_GRAPH: + if d != DEFAULT_GRAPH: + continue + elif g != GRAPH_WILDCARD and d != g: + continue + + results.append(QuadResult( + s=s, p=p, o=o, g=d, + otype=row.otype, dtype=row.dtype, lang=row.lang + )) + + return results + + def get_g(self, collection, g, limit=50): + """Get all quads in a specific graph""" + if g is None: + g = DEFAULT_GRAPH + + return self.session.execute(self.get_collection_by_graph_stmt, (collection, g, limit)) + + # ======================================================================== + # Collection management + # ======================================================================== + + def collection_exists(self, collection): + """Check if collection exists""" + try: + result = self.session.execute( + f"SELECT collection FROM {self.collection_metadata_table} WHERE collection = %s LIMIT 1", + (collection,) + ) + return bool(list(result)) + except Exception as e: + logger.error(f"Error checking collection existence: {e}") + return False + + def create_collection(self, collection): + """Create collection by inserting metadata row""" + try: + import datetime + self.session.execute( + f"INSERT INTO {self.collection_metadata_table} (collection, created_at) VALUES (%s, %s)", + (collection, datetime.datetime.now()) + ) + logger.info(f"Created collection metadata for {collection}") + except Exception as e: + logger.error(f"Error creating collection: {e}") + raise e + + def delete_collection(self, collection): + """ + Delete all quads for a collection from both tables. + + Uses efficient partition-level deletes: + 1. Read quads from quads_by_collection to get all quads + 2. Extract unique entities (s, p, o for URIs, g for non-default) + 3. Delete entire entity partitions + 4. Delete collection rows + """ + # Read all quads from collection table + rows = self.session.execute( + f"SELECT d, s, p, o, otype FROM {self.collection_table} WHERE collection = %s", + (collection,) + ) + + # Collect unique entities and quad data for deletion + entities = set() + quads = [] + + for row in rows: + d, s, p, o, otype = row.d, row.s, row.p, row.o, row.otype + quads.append((d, s, p, o)) + + # Subject and predicate are always entities + entities.add(s) + entities.add(p) + + # Object is an entity only for URIs + if otype == 'U' or otype == 'T': + entities.add(o) + + # Graph is an entity for non-default graphs + if d != DEFAULT_GRAPH: + entities.add(d) + + # Delete entity partitions (efficient partition-level deletes) + batch = BatchStatement() + count = 0 + + for entity in entities: + batch.add(self.delete_entity_partition_stmt, (collection, entity)) + count += 1 + + # Execute batch every 50 entities + if count % 50 == 0: + self.session.execute(batch) + batch = BatchStatement() + + # Execute remaining entity deletes + if count % 50 != 0: + self.session.execute(batch) + + # Delete collection rows + batch = BatchStatement() + count = 0 + + for d, s, p, o in quads: + batch.add(self.delete_collection_row_stmt, (collection, d, s, p, o)) + count += 1 + + # Execute batch every 50 quads + if count % 50 == 0: + self.session.execute(batch) + batch = BatchStatement() + + # Execute remaining collection row deletes + if count % 50 != 0: + self.session.execute(batch) + + # Delete collection metadata + self.session.execute( + f"DELETE FROM {self.collection_metadata_table} WHERE collection = %s", + (collection,) + ) + + logger.info(f"Deleted collection {collection}: {len(entities)} entity partitions, {len(quads)} quads") + + def close(self): + """Close connections""" + if hasattr(self, 'session') and self.session: + self.session.shutdown() + if hasattr(self, 'cluster') and self.cluster: + self.cluster.shutdown() + if self.cluster in _active_clusters: + _active_clusters.remove(self.cluster) + + +class QuadResult: + """ + Result object for quad queries, including object type metadata. + + Attributes: + s: Subject value + p: Predicate value + o: Object value + g: Graph/dataset value + otype: Object type - 'U' (URI), 'L' (literal), 'T' (triple) + dtype: XSD datatype (for literals) + lang: Language tag (for literals) + """ + + def __init__(self, s, p, o, g, otype='U', dtype='', lang=''): + self.s = s + self.p = p + self.o = o + self.g = g + self.otype = otype + self.dtype = dtype + self.lang = lang + + +def get_knowledge_graph_class(): + """ + Factory function to select KnowledgeGraph implementation. + + Uses CASSANDRA_ENTITY_CENTRIC environment variable to select: + - "true": EntityCentricKnowledgeGraph (new 2-table model) + - Otherwise: KnowledgeGraph (original 7-table model) + """ + use_entity_centric = os.environ.get("CASSANDRA_ENTITY_CENTRIC", "").lower() == "true" + if use_entity_centric: + logger.info("Using EntityCentricKnowledgeGraph (2-table model)") + return EntityCentricKnowledgeGraph + else: + logger.info("Using KnowledgeGraph (7-table model)") + return KnowledgeGraph diff --git a/trustgraph-flow/trustgraph/query/triples/cassandra/service.py b/trustgraph-flow/trustgraph/query/triples/cassandra/service.py index 51bbefa9..7b85a87c 100755 --- a/trustgraph-flow/trustgraph/query/triples/cassandra/service.py +++ b/trustgraph-flow/trustgraph/query/triples/cassandra/service.py @@ -6,7 +6,9 @@ null. Output is a list of quads. import logging -from .... direct.cassandra_kg import KnowledgeGraph, GRAPH_WILDCARD, DEFAULT_GRAPH +from .... direct.cassandra_kg import ( + KnowledgeGraph, GRAPH_WILDCARD, DEFAULT_GRAPH, get_knowledge_graph_class +) from .... schema import TriplesQueryRequest, TriplesQueryResponse, Error from .... schema import Term, Triple, IRI, LITERAL from .... base import TriplesQueryService @@ -31,8 +33,37 @@ def get_term_value(term): return term.id or term.value -def create_term(value): - """Create a Term from a string value""" +def create_term(value, otype=None, dtype=None, lang=None): + """ + Create a Term from a string value, optionally using type metadata. + + Args: + value: The string value + otype: Object type - 'U' (URI), 'L' (literal), 'T' (triple) + dtype: XSD datatype (for literals) + lang: Language tag (for literals) + + If otype is provided, uses it to determine Term type. + Otherwise falls back to URL detection heuristic. + """ + if otype is not None: + if otype == 'U': + return Term(type=IRI, iri=value) + elif otype == 'L': + return Term( + type=LITERAL, + value=value, + datatype=dtype or "", + language=lang or "" + ) + elif otype == 'T': + # Triple/reification - treat as IRI for now + return Term(type=IRI, iri=value) + else: + # Unknown otype, fall back to heuristic + pass + + # Heuristic fallback for backwards compatibility if value.startswith("http://") or value.startswith("https://"): return Term(type=IRI, iri=value) else: @@ -74,14 +105,17 @@ class Processor(TriplesQueryService): user = query.user if user != self.table: + # Use factory function to select implementation + KGClass = get_knowledge_graph_class() + if self.cassandra_username and self.cassandra_password: - self.tg = KnowledgeGraph( + self.tg = KGClass( hosts=self.cassandra_host, keyspace=query.user, username=self.cassandra_username, password=self.cassandra_password ) else: - self.tg = KnowledgeGraph( + self.tg = KGClass( hosts=self.cassandra_host, keyspace=query.user, ) @@ -93,6 +127,14 @@ class Processor(TriplesQueryService): o_val = get_term_value(query.o) g_val = query.g # Already a string or None + # Helper to extract object metadata from result row + def get_o_metadata(t): + """Extract otype/dtype/lang from result row if available""" + otype = getattr(t, 'otype', None) + dtype = getattr(t, 'dtype', None) + lang = getattr(t, 'lang', None) + return otype, dtype, lang + quads = [] # Route to appropriate query method based on which fields are specified @@ -106,7 +148,8 @@ class Processor(TriplesQueryService): ) for t in resp: g = t.g if hasattr(t, 'g') else DEFAULT_GRAPH - quads.append((s_val, p_val, o_val, g)) + otype, dtype, lang = get_o_metadata(t) + quads.append((s_val, p_val, o_val, g, otype, dtype, lang)) else: # SP specified resp = self.tg.get_sp( @@ -115,7 +158,8 @@ class Processor(TriplesQueryService): ) for t in resp: g = t.g if hasattr(t, 'g') else DEFAULT_GRAPH - quads.append((s_val, p_val, t.o, g)) + otype, dtype, lang = get_o_metadata(t) + quads.append((s_val, p_val, t.o, g, otype, dtype, lang)) else: if o_val is not None: # SO specified @@ -125,7 +169,8 @@ class Processor(TriplesQueryService): ) for t in resp: g = t.g if hasattr(t, 'g') else DEFAULT_GRAPH - quads.append((s_val, t.p, o_val, g)) + otype, dtype, lang = get_o_metadata(t) + quads.append((s_val, t.p, o_val, g, otype, dtype, lang)) else: # S only resp = self.tg.get_s( @@ -134,7 +179,8 @@ class Processor(TriplesQueryService): ) for t in resp: g = t.g if hasattr(t, 'g') else DEFAULT_GRAPH - quads.append((s_val, t.p, t.o, g)) + otype, dtype, lang = get_o_metadata(t) + quads.append((s_val, t.p, t.o, g, otype, dtype, lang)) else: if p_val is not None: if o_val is not None: @@ -145,7 +191,8 @@ class Processor(TriplesQueryService): ) for t in resp: g = t.g if hasattr(t, 'g') else DEFAULT_GRAPH - quads.append((t.s, p_val, o_val, g)) + otype, dtype, lang = get_o_metadata(t) + quads.append((t.s, p_val, o_val, g, otype, dtype, lang)) else: # P only resp = self.tg.get_p( @@ -154,7 +201,8 @@ class Processor(TriplesQueryService): ) for t in resp: g = t.g if hasattr(t, 'g') else DEFAULT_GRAPH - quads.append((t.s, p_val, t.o, g)) + otype, dtype, lang = get_o_metadata(t) + quads.append((t.s, p_val, t.o, g, otype, dtype, lang)) else: if o_val is not None: # O only @@ -164,7 +212,8 @@ class Processor(TriplesQueryService): ) for t in resp: g = t.g if hasattr(t, 'g') else DEFAULT_GRAPH - quads.append((t.s, t.p, o_val, g)) + otype, dtype, lang = get_o_metadata(t) + quads.append((t.s, t.p, o_val, g, otype, dtype, lang)) else: # Nothing specified - get all resp = self.tg.get_all( @@ -173,14 +222,16 @@ class Processor(TriplesQueryService): ) for t in resp: g = t.g if hasattr(t, 'g') else DEFAULT_GRAPH - quads.append((t.s, t.p, t.o, g)) + otype, dtype, lang = get_o_metadata(t) + quads.append((t.s, t.p, t.o, g, otype, dtype, lang)) # Convert to Triple objects (with g field) + # Use otype/dtype/lang for proper Term reconstruction if available triples = [ Triple( s=create_term(q[0]), p=create_term(q[1]), - o=create_term(q[2]), + o=create_term(q[2], otype=q[4], dtype=q[5], lang=q[6]), g=q[3] if q[3] != DEFAULT_GRAPH else None ) for q in quads diff --git a/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py b/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py index 6ed93b7d..5d1e160b 100755 --- a/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py +++ b/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py @@ -10,12 +10,14 @@ import argparse import time import logging -from .... direct.cassandra_kg import KnowledgeGraph, DEFAULT_GRAPH +from .... direct.cassandra_kg import ( + KnowledgeGraph, DEFAULT_GRAPH, get_knowledge_graph_class +) from .... base import TriplesStoreService, CollectionConfigHandler from .... base import AsyncProcessor, Consumer, Producer from .... base import ConsumerMetrics, ProducerMetrics from .... base.cassandra_config import add_cassandra_args, resolve_cassandra_config -from .... schema import IRI, LITERAL +from .... schema import IRI, LITERAL, BLANK, TRIPLE # Module logger logger = logging.getLogger(__name__) @@ -36,6 +38,46 @@ def get_term_value(term): return term.id or term.value +def get_term_otype(term): + """ + Get object type code from a Term for entity-centric storage. + + Maps Term.type to otype: + - IRI ("i") → "U" (URI) + - BLANK ("b") → "U" (treated as URI) + - LITERAL ("l") → "L" (Literal) + - TRIPLE ("t") → "T" (Triple/reification) + """ + if term is None: + return "U" + if term.type == IRI or term.type == BLANK: + return "U" + elif term.type == LITERAL: + return "L" + elif term.type == TRIPLE: + return "T" + else: + return "U" + + +def get_term_dtype(term): + """Extract datatype from a Term (for literals)""" + if term is None: + return "" + if term.type == LITERAL: + return term.datatype or "" + return "" + + +def get_term_lang(term): + """Extract language tag from a Term (for literals)""" + if term is None: + return "" + if term.type == LITERAL: + return term.language or "" + return "" + + class Processor(CollectionConfigHandler, TriplesStoreService): def __init__(self, **params): @@ -78,15 +120,18 @@ class Processor(CollectionConfigHandler, TriplesStoreService): self.tg = None + # Use factory function to select implementation + KGClass = get_knowledge_graph_class() + try: if self.cassandra_username and self.cassandra_password: - self.tg = KnowledgeGraph( + self.tg = KGClass( hosts=self.cassandra_host, keyspace=message.metadata.user, username=self.cassandra_username, password=self.cassandra_password ) else: - self.tg = KnowledgeGraph( + self.tg = KGClass( hosts=self.cassandra_host, keyspace=message.metadata.user, ) @@ -105,12 +150,20 @@ class Processor(CollectionConfigHandler, TriplesStoreService): # t.g is None for default graph, or a graph IRI g_val = t.g if t.g is not None else DEFAULT_GRAPH + # Extract object type metadata for entity-centric storage + otype = get_term_otype(t.o) + dtype = get_term_dtype(t.o) + lang = get_term_lang(t.o) + self.tg.insert( message.metadata.collection, s_val, p_val, o_val, - g=g_val + g=g_val, + otype=otype, + dtype=dtype, + lang=lang ) async def create_collection(self, user: str, collection: str, metadata: dict): @@ -120,16 +173,19 @@ class Processor(CollectionConfigHandler, TriplesStoreService): if self.table is None or self.table != user: self.tg = None + # Use factory function to select implementation + KGClass = get_knowledge_graph_class() + try: if self.cassandra_username and self.cassandra_password: - self.tg = KnowledgeGraph( + self.tg = KGClass( hosts=self.cassandra_host, keyspace=user, username=self.cassandra_username, password=self.cassandra_password ) else: - self.tg = KnowledgeGraph( + self.tg = KGClass( hosts=self.cassandra_host, keyspace=user, ) @@ -159,16 +215,19 @@ class Processor(CollectionConfigHandler, TriplesStoreService): if self.table is None or self.table != user: self.tg = None + # Use factory function to select implementation + KGClass = get_knowledge_graph_class() + try: if self.cassandra_username and self.cassandra_password: - self.tg = KnowledgeGraph( + self.tg = KGClass( hosts=self.cassandra_host, keyspace=user, username=self.cassandra_username, password=self.cassandra_password ) else: - self.tg = KnowledgeGraph( + self.tg = KGClass( hosts=self.cassandra_host, keyspace=user, )