wip: add user defined types

This commit is contained in:
Namekuji 2023-09-07 12:16:37 -04:00
parent ede4e436f6
commit 8b46d1de03
No known key found for this signature in database
GPG key ID: 1D62332C07FBA532
13 changed files with 562 additions and 5 deletions

View file

@ -2445,6 +2445,7 @@ version = "0.1.0"
dependencies = [ dependencies = [
"chrono", "chrono",
"clap", "clap",
"futures",
"scylla", "scylla",
"sea-orm", "sea-orm",
"serde", "serde",

View file

@ -8,6 +8,7 @@ edition = "2021"
[dependencies] [dependencies]
chrono = "0.4.26" chrono = "0.4.26"
clap = { version = "4.3.11", features = ["derive"] } clap = { version = "4.3.11", features = ["derive"] }
futures = "0.3.28"
scylla = "0.8.2" scylla = "0.8.2"
sea-orm = { version = "0.12.2", features = ["sqlx-postgres", "runtime-tokio-rustls"] } sea-orm = { version = "0.12.2", features = ["sqlx-postgres", "runtime-tokio-rustls"] }
serde = { version = "1.0.171", features = ["derive"] } serde = { version = "1.0.171", features = ["derive"] }

View file

@ -0,0 +1,60 @@
//! `SeaORM` Entity. Generated by sea-orm-codegen 0.12.2
use sea_orm::entity::prelude::*;
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)]
#[sea_orm(table_name = "drive_file")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
#[sea_orm(column_name = "createdAt")]
pub created_at: DateTimeWithTimeZone,
#[sea_orm(column_name = "userId")]
pub user_id: Option<String>,
#[sea_orm(column_name = "userHost")]
pub user_host: Option<String>,
pub md5: String,
pub name: String,
pub r#type: String,
pub size: i32,
pub comment: Option<String>,
#[sea_orm(column_type = "JsonBinary")]
pub properties: Json,
#[sea_orm(column_name = "storedInternal")]
pub stored_internal: bool,
pub url: String,
#[sea_orm(column_name = "thumbnailUrl")]
pub thumbnail_url: Option<String>,
#[sea_orm(column_name = "webpublicUrl")]
pub webpublic_url: Option<String>,
#[sea_orm(column_name = "accessKey")]
pub access_key: Option<String>,
#[sea_orm(column_name = "thumbnailAccessKey")]
pub thumbnail_access_key: Option<String>,
#[sea_orm(column_name = "webpublicAccessKey")]
pub webpublic_access_key: Option<String>,
pub uri: Option<String>,
pub src: Option<String>,
#[sea_orm(column_name = "folderId")]
pub folder_id: Option<String>,
#[sea_orm(column_name = "isSensitive")]
pub is_sensitive: bool,
#[sea_orm(column_name = "isLink")]
pub is_link: bool,
pub blurhash: Option<String>,
#[sea_orm(column_name = "webpublicType")]
pub webpublic_type: Option<String>,
#[sea_orm(column_name = "requestHeaders", column_type = "JsonBinary", nullable)]
pub request_headers: Option<Json>,
#[sea_orm(column_name = "requestIp")]
pub request_ip: Option<String>,
#[sea_orm(column_name = "maybeSensitive")]
pub maybe_sensitive: bool,
#[sea_orm(column_name = "maybePorn")]
pub maybe_porn: bool,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}

View file

@ -0,0 +1,30 @@
//! `SeaORM` Entity. Generated by sea-orm-codegen 0.12.2
use sea_orm::entity::prelude::*;
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)]
#[sea_orm(table_name = "emoji")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
#[sea_orm(column_name = "updatedAt")]
pub updated_at: Option<DateTimeWithTimeZone>,
pub name: String,
pub host: Option<String>,
#[sea_orm(column_name = "originalUrl")]
pub original_url: String,
pub uri: Option<String>,
pub r#type: Option<String>,
pub aliases: Vec<String>,
pub category: Option<String>,
#[sea_orm(column_name = "publicUrl")]
pub public_url: String,
pub license: Option<String>,
pub width: Option<i32>,
pub height: Option<i32>,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}

View file

@ -0,0 +1,6 @@
pub mod drive_file;
pub mod emoji;
pub mod note;
pub mod poll_vote;
pub mod poll;
pub mod sea_orm_active_enums;

View file

@ -0,0 +1,69 @@
//! `SeaORM` Entity. Generated by sea-orm-codegen 0.12.2
use super::sea_orm_active_enums::NoteVisibilityEnum;
use sea_orm::entity::prelude::*;
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)]
#[sea_orm(table_name = "note")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
#[sea_orm(column_name = "createdAt")]
pub created_at: DateTimeWithTimeZone,
#[sea_orm(column_name = "replyId")]
pub reply_id: Option<String>,
#[sea_orm(column_name = "renoteId")]
pub renote_id: Option<String>,
#[sea_orm(column_type = "Text", nullable)]
pub text: Option<String>,
pub name: Option<String>,
pub cw: Option<String>,
#[sea_orm(column_name = "userId")]
pub user_id: String,
#[sea_orm(column_name = "localOnly")]
pub local_only: bool,
#[sea_orm(column_name = "renoteCount")]
pub renote_count: i16,
#[sea_orm(column_name = "repliesCount")]
pub replies_count: i16,
#[sea_orm(column_type = "JsonBinary")]
pub reactions: Json,
pub visibility: NoteVisibilityEnum,
pub uri: Option<String>,
pub score: i32,
#[sea_orm(column_name = "fileIds")]
pub file_ids: Vec<String>,
#[sea_orm(column_name = "attachedFileTypes")]
pub attached_file_types: Vec<String>,
#[sea_orm(column_name = "visibleUserIds")]
pub visible_user_ids: Vec<String>,
pub mentions: Vec<String>,
#[sea_orm(column_name = "mentionedRemoteUsers", column_type = "Text")]
pub mentioned_remote_users: String,
pub emojis: Vec<String>,
pub tags: Vec<String>,
#[sea_orm(column_name = "hasPoll")]
pub has_poll: bool,
#[sea_orm(column_name = "userHost")]
pub user_host: Option<String>,
#[sea_orm(column_name = "replyUserId")]
pub reply_user_id: Option<String>,
#[sea_orm(column_name = "replyUserHost")]
pub reply_user_host: Option<String>,
#[sea_orm(column_name = "renoteUserId")]
pub renote_user_id: Option<String>,
#[sea_orm(column_name = "renoteUserHost")]
pub renote_user_host: Option<String>,
pub url: Option<String>,
#[sea_orm(column_name = "channelId")]
pub channel_id: Option<String>,
#[sea_orm(column_name = "threadId")]
pub thread_id: Option<String>,
#[sea_orm(column_name = "updatedAt")]
pub updated_at: Option<DateTimeWithTimeZone>,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}

View file

@ -0,0 +1,42 @@
//! `SeaORM` Entity. Generated by sea-orm-codegen 0.12.2
use super::sea_orm_active_enums::PollNotevisibilityEnum;
use sea_orm::entity::prelude::*;
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)]
#[sea_orm(table_name = "poll")]
pub struct Model {
#[sea_orm(column_name = "noteId", primary_key, auto_increment = false, unique)]
pub note_id: String,
#[sea_orm(column_name = "expiresAt")]
pub expires_at: Option<DateTimeWithTimeZone>,
pub multiple: bool,
pub choices: Vec<String>,
pub votes: Vec<i32>,
#[sea_orm(column_name = "noteVisibility")]
pub note_visibility: PollNotevisibilityEnum,
#[sea_orm(column_name = "userId")]
pub user_id: String,
#[sea_orm(column_name = "userHost")]
pub user_host: Option<String>,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(
belongs_to = "super::note::Entity",
from = "Column::NoteId",
to = "super::note::Column::Id",
on_update = "NoAction",
on_delete = "Cascade"
)]
Note,
}
impl Related<super::note::Entity> for Entity {
fn to() -> RelationDef {
Relation::Note.def()
}
}
impl ActiveModelBehavior for ActiveModel {}

View file

@ -0,0 +1,37 @@
//! `SeaORM` Entity. Generated by sea-orm-codegen 0.12.2
use sea_orm::entity::prelude::*;
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)]
#[sea_orm(table_name = "poll_vote")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub id: String,
#[sea_orm(column_name = "createdAt")]
pub created_at: DateTimeWithTimeZone,
#[sea_orm(column_name = "userId")]
pub user_id: String,
#[sea_orm(column_name = "noteId")]
pub note_id: String,
pub choice: i32,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(
belongs_to = "super::note::Entity",
from = "Column::NoteId",
to = "super::note::Column::Id",
on_update = "NoAction",
on_delete = "Cascade"
)]
Note,
}
impl Related<super::note::Entity> for Entity {
fn to() -> RelationDef {
Relation::Note.def()
}
}
impl ActiveModelBehavior for ActiveModel {}

View file

@ -0,0 +1,206 @@
//! `SeaORM` Entity. Generated by sea-orm-codegen 0.12.2
use sea_orm::entity::prelude::*;
#[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum)]
#[sea_orm(rs_type = "String", db_type = "Enum", enum_name = "antenna_src_enum")]
pub enum AntennaSrcEnum {
#[sea_orm(string_value = "all")]
All,
#[sea_orm(string_value = "group")]
Group,
#[sea_orm(string_value = "home")]
Home,
#[sea_orm(string_value = "instances")]
Instances,
#[sea_orm(string_value = "list")]
List,
#[sea_orm(string_value = "users")]
Users,
}
#[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum)]
#[sea_orm(
rs_type = "String",
db_type = "Enum",
enum_name = "meta_sensitivemediadetection_enum"
)]
pub enum MetaSensitivemediadetectionEnum {
#[sea_orm(string_value = "all")]
All,
#[sea_orm(string_value = "local")]
Local,
#[sea_orm(string_value = "none")]
None,
#[sea_orm(string_value = "remote")]
Remote,
}
#[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum)]
#[sea_orm(
rs_type = "String",
db_type = "Enum",
enum_name = "meta_sensitivemediadetectionsensitivity_enum"
)]
pub enum MetaSensitivemediadetectionsensitivityEnum {
#[sea_orm(string_value = "high")]
High,
#[sea_orm(string_value = "low")]
Low,
#[sea_orm(string_value = "medium")]
Medium,
#[sea_orm(string_value = "veryHigh")]
VeryHigh,
#[sea_orm(string_value = "veryLow")]
VeryLow,
}
#[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum)]
#[sea_orm(
rs_type = "String",
db_type = "Enum",
enum_name = "muted_note_reason_enum"
)]
pub enum MutedNoteReasonEnum {
#[sea_orm(string_value = "manual")]
Manual,
#[sea_orm(string_value = "other")]
Other,
#[sea_orm(string_value = "spam")]
Spam,
#[sea_orm(string_value = "word")]
Word,
}
#[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum)]
#[sea_orm(
rs_type = "String",
db_type = "Enum",
enum_name = "note_visibility_enum"
)]
pub enum NoteVisibilityEnum {
#[sea_orm(string_value = "followers")]
Followers,
#[sea_orm(string_value = "hidden")]
Hidden,
#[sea_orm(string_value = "home")]
Home,
#[sea_orm(string_value = "public")]
Public,
#[sea_orm(string_value = "specified")]
Specified,
}
#[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum)]
#[sea_orm(
rs_type = "String",
db_type = "Enum",
enum_name = "notification_type_enum"
)]
pub enum NotificationTypeEnum {
#[sea_orm(string_value = "app")]
App,
#[sea_orm(string_value = "follow")]
Follow,
#[sea_orm(string_value = "followRequestAccepted")]
FollowRequestAccepted,
#[sea_orm(string_value = "groupInvited")]
GroupInvited,
#[sea_orm(string_value = "mention")]
Mention,
#[sea_orm(string_value = "pollEnded")]
PollEnded,
#[sea_orm(string_value = "pollVote")]
PollVote,
#[sea_orm(string_value = "quote")]
Quote,
#[sea_orm(string_value = "reaction")]
Reaction,
#[sea_orm(string_value = "receiveFollowRequest")]
ReceiveFollowRequest,
#[sea_orm(string_value = "renote")]
Renote,
#[sea_orm(string_value = "reply")]
Reply,
}
#[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum)]
#[sea_orm(
rs_type = "String",
db_type = "Enum",
enum_name = "page_visibility_enum"
)]
pub enum PageVisibilityEnum {
#[sea_orm(string_value = "followers")]
Followers,
#[sea_orm(string_value = "public")]
Public,
#[sea_orm(string_value = "specified")]
Specified,
}
#[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum)]
#[sea_orm(
rs_type = "String",
db_type = "Enum",
enum_name = "poll_notevisibility_enum"
)]
pub enum PollNotevisibilityEnum {
#[sea_orm(string_value = "followers")]
Followers,
#[sea_orm(string_value = "home")]
Home,
#[sea_orm(string_value = "public")]
Public,
#[sea_orm(string_value = "specified")]
Specified,
}
#[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum)]
#[sea_orm(rs_type = "String", db_type = "Enum", enum_name = "relay_status_enum")]
pub enum RelayStatusEnum {
#[sea_orm(string_value = "accepted")]
Accepted,
#[sea_orm(string_value = "rejected")]
Rejected,
#[sea_orm(string_value = "requesting")]
Requesting,
}
#[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum)]
#[sea_orm(
rs_type = "String",
db_type = "Enum",
enum_name = "user_profile_ffvisibility_enum"
)]
pub enum UserProfileFfvisibilityEnum {
#[sea_orm(string_value = "followers")]
Followers,
#[sea_orm(string_value = "private")]
Private,
#[sea_orm(string_value = "public")]
Public,
}
#[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum)]
#[sea_orm(
rs_type = "String",
db_type = "Enum",
enum_name = "user_profile_mutingnotificationtypes_enum"
)]
pub enum UserProfileMutingnotificationtypesEnum {
#[sea_orm(string_value = "app")]
App,
#[sea_orm(string_value = "follow")]
Follow,
#[sea_orm(string_value = "followRequestAccepted")]
FollowRequestAccepted,
#[sea_orm(string_value = "groupInvited")]
GroupInvited,
#[sea_orm(string_value = "mention")]
Mention,
#[sea_orm(string_value = "pollEnded")]
PollEnded,
#[sea_orm(string_value = "pollVote")]
PollVote,
#[sea_orm(string_value = "quote")]
Quote,
#[sea_orm(string_value = "reaction")]
Reaction,
#[sea_orm(string_value = "receiveFollowRequest")]
ReceiveFollowRequest,
#[sea_orm(string_value = "renote")]
Renote,
#[sea_orm(string_value = "reply")]
Reply,
}

View file

@ -3,6 +3,7 @@ pub use clap::Parser;
pub mod cli; pub mod cli;
pub mod config; pub mod config;
pub mod error; pub mod error;
pub mod entity;
pub(crate) mod migrator; pub(crate) mod migrator;
pub(crate) mod setup; pub(crate) mod setup;

View file

@ -1,9 +1,15 @@
use scylla::{Session, SessionBuilder}; use std::collections::HashMap;
use sea_orm::{ConnectionTrait, Database, Statement};
use futures::TryStreamExt;
use scylla::{Session, SessionBuilder, IntoUserType, FromUserType, ValueList};
use sea_orm::{ConnectionTrait, Database, EntityTrait, Statement};
use urlencoding::encode; use urlencoding::encode;
use chrono::{DateTime, Utc, NaiveDate};
use crate::{ use crate::{
config::{DbConfig, ScyllaConfig}, config::{DbConfig, ScyllaConfig},
entity::note,
error::Error, error::Error,
}; };
@ -41,6 +47,11 @@ impl Initializer {
let pool = Database::connect(&self.postgres_url).await?; let pool = Database::connect(&self.postgres_url).await?;
let db_backend = pool.get_database_backend(); let db_backend = pool.get_database_backend();
let mut notes = note::Entity::find().stream(&pool).await?;
while let Some(note) = notes.try_next().await? {
}
let fk_pairs = vec![ let fk_pairs = vec![
("channel_note_pining", "FK_10b19ef67d297ea9de325cd4502"), ("channel_note_pining", "FK_10b19ef67d297ea9de325cd4502"),
("clip_note", "FK_a012eaf5c87c65da1deb5fdbfa3"), ("clip_note", "FK_a012eaf5c87c65da1deb5fdbfa3"),
@ -79,3 +90,86 @@ impl Initializer {
Ok(()) Ok(())
} }
} }
#[derive(Debug, IntoUserType, FromUserType)]
struct DriveFileType {
id: String,
#[scylla_crate(rename = "type")]
kind: String,
#[scylla_crate(rename = "createdAt")]
created_at: DateTime<Utc>,
name: String,
comment: Option<String>,
blurhash: Option<String>,
url: String,
#[scylla_crate(rename = "thumbnailUrl")]
thumbnail_url: Option<String>,
#[scylla_crate(rename = "isSensitive")]
is_sensitive: bool,
#[scylla_crate(rename = "isLink")]
is_link: bool,
md5: String,
size: i32,
width: Option<i32>,
height: Option<i32>,
}
#[derive(Debug, IntoUserType, FromUserType)]
struct NoteEditHistoryType {
content: Option<String>,
cw: Option<String>,
files: Vec<DriveFileType>,
#[scylla_crate(rename = "updatedAt")]
updated_at: DateTime<Utc>
}
#[derive(Debug, IntoUserType, FromUserType)]
struct EmojiType {
name: String,
url: String,
width: Option<i32>,
height: Option<i32>,
}
#[derive(Debug, IntoUserType, FromUserType)]
struct PollType {
#[scylla_crate(rename = "expiresAt")]
expires_at: Option<DateTime<Utc>>,
multiple: bool,
choices: HashMap<i32, String>,
}
#[derive(ValueList)]
struct NoteTable {
#[scylla_crate(rename = "createdAtDate")]
created_at_date: NaiveDate,
#[scylla_crate(rename = "createdAt")]
created_at: DateTime<Utc>,
id: String,
visibility: String,
content: Option<String>,
name: Option<String>,
cw: Option<String>,
#[scylla_crate(rename = "localOnly")]
local_only: bool,
#[scylla_crate(rename = "renoteCount")]
renote_count: i32,
#[scylla_crate(rename = "scyllaCrate")]
replies_count: i32,
uri: Option<String>,
url: Option<String>,
score: i32,
files: Vec<DriveFileType>,
#[scylla_crate(rename = "visibleUserIds")]
visible_user_ids: Vec<String>,
mentions: Vec<String>,
#[scylla_crate(rename = "mentionedRemoteUsers")]
mentioned_remote_users: String,
emojis: Vec<String>,
tags: Vec<String>,
#[scylla_crate(rename = "hasPoll")]
has_poll: bool,
poll: PollType,
#[scylla_crate(rename = "threadId")]
thread_id: String,
}

View file

@ -32,10 +32,17 @@ function newClient(): Client | null {
const requestTracker = new tracker.RequestLogger({ const requestTracker = new tracker.RequestLogger({
slowThreshold: 1000, slowThreshold: 1000,
}); });
const client = new Client({ const client = new Client({
contactPoints: config.scylla.nodes, contactPoints: config.scylla.nodes,
localDataCenter: config.scylla.localDataCentre, localDataCenter: config.scylla.localDataCentre,
keyspace: config.scylla.keyspace, keyspace: config.scylla.keyspace,
pooling: {
coreConnectionsPerHost: {
[types.distance.local]: 2,
[types.distance.remote]: 1,
},
},
requestTracker, requestTracker,
}); });
@ -79,14 +86,16 @@ export async function fetchPostCount(local = false): Promise<number> {
if (local) { if (local) {
return await localPostCountCache.fetch(null, () => return await localPostCountCache.fetch(null, () =>
scyllaClient scyllaClient
.execute("SELECT COUNT(*) FROM local_note_by_user_id") .execute("SELECT COUNT(*) FROM local_note_by_user_id", [], {
prepare: true,
})
.then((result) => result.first().get("count") as number), .then((result) => result.first().get("count") as number),
); );
} }
return await allPostCountCache.fetch(null, () => return await allPostCountCache.fetch(null, () =>
scyllaClient scyllaClient
.execute("SELECT COUNT(*) FROM note") .execute("SELECT COUNT(*) FROM note", [], { prepare: true })
.then((result) => result.first().get("count") as number), .then((result) => result.first().get("count") as number),
); );
} }
@ -98,7 +107,7 @@ export async function fetchReactionCount(): Promise<number> {
return await reactionCountCache.fetch(null, () => return await reactionCountCache.fetch(null, () =>
scyllaClient scyllaClient
.execute("SELECT COUNT(*) FROM reaction") .execute("SELECT COUNT(*) FROM reaction", [], { prepare: true })
.then((result) => result.first().get("count") as number), .then((result) => result.first().get("count") as number),
); );
} }

View file

@ -26,6 +26,7 @@ export default async function indexAllNotes(
scyllaClient.eachRow( scyllaClient.eachRow(
"SELECT * FROM note", "SELECT * FROM note",
[], [],
{ prepare: true },
(n, row) => { (n, row) => {
if (n % 1000 === 0) { if (n % 1000 === 0) {
job.progress(((n / total) * 100).toFixed(1)); job.progress(((n / total) * 100).toFixed(1));