matrix-rust-sdk

public

Star0

Watch0

61 commits 4 branches 0 tags 0 releases16.6 MiB

// Copyright 2024 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use ruma::events::room::message::{MessageType, OriginalSyncRoomMessageEvent, Relation};
use tantivy::{
    DateTime, TantivyDocument, doc,
    schema::{DateOptions, DateTimePrecision, Field, INDEXED, STORED, STRING, Schema, TEXT},
};

use crate::error::{IndexError, IndexSchemaError};

pub(crate) trait MatrixSearchIndexSchema {
    fn new() -> Self;
    fn default_search_fields(&self) -> Vec<Field>;
    fn primary_key(&self) -> Field;
    fn deletion_key(&self) -> Field;
    fn get_field_name(&self, field: Field) -> &str;
    fn as_tantivy_schema(&self) -> Schema;
    fn make_doc(&self, event: OriginalSyncRoomMessageEvent) -> Result<TantivyDocument, IndexError>;
}

#[derive(Debug, Clone)]
pub(crate) struct RoomMessageSchema {
    inner: Schema,
    /// The event id of this event (primary key).
    event_id_field: Field,
    /// The event id of the event that this event affects.
    /// Used by edits to refer to the event they edited (deletion key).
    original_event_id_field: Field,
    body_field: Field,
    date_field: Field,
    sender_field: Field,
    default_search_fields: Vec<Field>,
}

impl MatrixSearchIndexSchema for RoomMessageSchema {
    fn new() -> Self {
        let mut schema = Schema::builder();
        let event_id_field = schema.add_text_field("event_id", STORED | STRING);
        let original_event_id_field = schema.add_text_field("original_event_id", STRING);
        let body_field = schema.add_text_field("body", TEXT);

        let date_options =
            DateOptions::from(INDEXED).set_fast().set_precision(DateTimePrecision::Seconds);

        let date_field = schema.add_date_field("date", date_options);
        let sender_field = schema.add_text_field("sender", STRING);

        let default_search_fields = vec![body_field];

        let schema = schema.build();

        Self {
            inner: schema,
            event_id_field,
            original_event_id_field,
            body_field,
            date_field,
            sender_field,
            default_search_fields,
        }
    }

    fn default_search_fields(&self) -> Vec<Field> {
        self.default_search_fields.clone()
    }

    fn primary_key(&self) -> Field {
        self.event_id_field
    }

    fn deletion_key(&self) -> Field {
        self.original_event_id_field
    }

    fn get_field_name(&self, field: Field) -> &str {
        self.inner.get_field_name(field)
    }

    fn as_tantivy_schema(&self) -> Schema {
        self.inner.clone()
    }

    /// Given an [`OriginalSyncRoomMessageEvent`] return a
    /// [`TantivyDocument`].
    fn make_doc(&self, event: OriginalSyncRoomMessageEvent) -> Result<TantivyDocument, IndexError> {
        let body = match &event.content.msgtype {
            MessageType::Text(content) => Ok(content.body.clone()),
            _ => Err(IndexError::MessageTypeNotSupported),
        }?;

        let mut document = doc!(
            self.event_id_field => event.event_id.to_string(),
            self.body_field => body,
            self.date_field =>
                DateTime::from_timestamp_millis(
                    event.origin_server_ts.get().into()),
            self.sender_field => event.sender.to_string(),
        );

        if let Some(Relation::Replacement(replacement_data)) = &event.content.relates_to {
            document.add_text(self.original_event_id_field, replacement_data.event_id.clone());
        } else {
            document.add_text(self.original_event_id_field, event.event_id);
        }

        Ok(document)
    }
}

impl TryFrom<Schema> for RoomMessageSchema {
    type Error = IndexSchemaError;

    fn try_from(schema: Schema) -> Result<RoomMessageSchema, Self::Error> {
        let event_id_field = schema.get_field("event_id")?;
        let original_event_id_field = schema.get_field("original_event_id")?;
        let body_field = schema.get_field("body")?;
        let date_field = schema.get_field("date")?;
        let sender_field = schema.get_field("sender")?;

        let default_search_fields = vec![body_field];

        Ok(Self {
            inner: schema,
            event_id_field,
            original_event_id_field,
            body_field,
            date_field,
            sender_field,
            default_search_fields,
        })
    }
}

1	// Copyright 2024 The Matrix.org Foundation C.I.C.
2	//
3	// Licensed under the Apache License, Version 2.0 (the "License");
4	// you may not use this file except in compliance with the License.
5	// You may obtain a copy of the License at
6	//
7	// http://www.apache.org/licenses/LICENSE-2.0
8	//
9	// Unless required by applicable law or agreed to in writing, software
10	// distributed under the License is distributed on an "AS IS" BASIS,
11	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12	// See the License for the specific language governing permissions and
13	// limitations under the License.
14
15	use ruma::events::room::message::{MessageType, OriginalSyncRoomMessageEvent, Relation};
16	use tantivy::{
17	DateTime, TantivyDocument, doc,
18	schema::{DateOptions, DateTimePrecision, Field, INDEXED, STORED, STRING, Schema, TEXT},
19	};
20
21	use crate::error::{IndexError, IndexSchemaError};
22
23	pub(crate) trait MatrixSearchIndexSchema {
24	fn new() -> Self;
25	fn default_search_fields(&self) -> Vec<Field>;
26	fn primary_key(&self) -> Field;
27	fn deletion_key(&self) -> Field;
28	fn get_field_name(&self, field: Field) -> &str;
29	fn as_tantivy_schema(&self) -> Schema;
30	fn make_doc(&self, event: OriginalSyncRoomMessageEvent) -> Result<TantivyDocument, IndexError>;
31	}
32
33	#[derive(Debug, Clone)]
34	pub(crate) struct RoomMessageSchema {
35	inner: Schema,
36	/// The event id of this event (primary key).
37	event_id_field: Field,
38	/// The event id of the event that this event affects.
39	/// Used by edits to refer to the event they edited (deletion key).
40	original_event_id_field: Field,
41	body_field: Field,
42	date_field: Field,
43	sender_field: Field,
44	default_search_fields: Vec<Field>,
45	}
46
47	impl MatrixSearchIndexSchema for RoomMessageSchema {
48	fn new() -> Self {
49	let mut schema = Schema::builder();
50	let event_id_field = schema.add_text_field("event_id", STORED \| STRING);
51	let original_event_id_field = schema.add_text_field("original_event_id", STRING);
52	let body_field = schema.add_text_field("body", TEXT);
53
54	let date_options =
55	DateOptions::from(INDEXED).set_fast().set_precision(DateTimePrecision::Seconds);
56
57	let date_field = schema.add_date_field("date", date_options);
58	let sender_field = schema.add_text_field("sender", STRING);
59
60	let default_search_fields = vec![body_field];
61
62	let schema = schema.build();
63
64	Self {
65	inner: schema,
66	event_id_field,
67	original_event_id_field,
68	body_field,
69	date_field,
70	sender_field,
71	default_search_fields,
72	}
73	}
74
75	fn default_search_fields(&self) -> Vec<Field> {
76	self.default_search_fields.clone()
77	}
78
79	fn primary_key(&self) -> Field {
80	self.event_id_field
81	}
82
83	fn deletion_key(&self) -> Field {
84	self.original_event_id_field
85	}
86
87	fn get_field_name(&self, field: Field) -> &str {
88	self.inner.get_field_name(field)
89	}
90
91	fn as_tantivy_schema(&self) -> Schema {
92	self.inner.clone()
93	}
94
95	/// Given an [`OriginalSyncRoomMessageEvent`] return a
96	/// [`TantivyDocument`].
97	fn make_doc(&self, event: OriginalSyncRoomMessageEvent) -> Result<TantivyDocument, IndexError> {
98	let body = match &event.content.msgtype {
99	MessageType::Text(content) => Ok(content.body.clone()),
100	_ => Err(IndexError::MessageTypeNotSupported),
101	}?;
102
103	let mut document = doc!(
104	self.event_id_field => event.event_id.to_string(),
105	self.body_field => body,
106	self.date_field =>
107	DateTime::from_timestamp_millis(
108	event.origin_server_ts.get().into()),
109	self.sender_field => event.sender.to_string(),
110	);
111
112	if let Some(Relation::Replacement(replacement_data)) = &event.content.relates_to {
113	document.add_text(self.original_event_id_field, replacement_data.event_id.clone());
114	} else {
115	document.add_text(self.original_event_id_field, event.event_id);
116	}
117
118	Ok(document)
119	}
120	}
121
122	impl TryFrom<Schema> for RoomMessageSchema {
123	type Error = IndexSchemaError;
124
125	fn try_from(schema: Schema) -> Result<RoomMessageSchema, Self::Error> {
126	let event_id_field = schema.get_field("event_id")?;
127	let original_event_id_field = schema.get_field("original_event_id")?;
128	let body_field = schema.get_field("body")?;
129	let date_field = schema.get_field("date")?;
130	let sender_field = schema.get_field("sender")?;
131
132	let default_search_fields = vec![body_field];
133
134	Ok(Self {
135	inner: schema,
136	event_id_field,
137	original_event_id_field,
138	body_field,
139	date_field,
140	sender_field,
141	default_search_fields,
142	})
143	}
144	}