Function get_samples

tests/graph/benches/representative_read/seed.rs:293–362

(account_id: ActorEntityUuid, store_wrapper: &StoreWrapper)

}

async fn get_samples(account_id: ActorEntityUuid, store_wrapper: &StoreWrapper) -> Samples {
    let mut entity_types = HashMap::new();
    entity_types.insert(
        account_id,
        SEED_ENTITY_TYPES
            .into_iter()
            .map(|entity_type_str| {
                let entity_type: EntityType =
                    serde_json::from_str(entity_type_str).expect("could not parse entity type");
                entity_type.id
            })
            .collect(),
    );

    let mut samples = Samples {
        entities: HashMap::from([(account_id, HashMap::new())]),
        entity_types,
    };

    let sample_map = samples
        .entities
        .get_mut(&account_id)
        .expect("could not get sample map");

    for entity_type_id in SEED_ENTITIES.map(|(entity_type_str, ..)| {
        let entity_type: EntityType =
            serde_json::from_str(entity_type_str).expect("could not parse entity type");
        entity_type.id
    }) {
        // For now we'll just pick a sample of 50 entities.
        let sample_entity_uuids = store_wrapper
            .store
            .as_client()
            .query(
                "
                -- Very naive and slow sampling, we can replace when this becomes a bottleneck
                SELECT entity_uuid FROM entity_temporal_metadata
                INNER JOIN entity_is_of_type ON entity_is_of_type.entity_edition_id = \
                 entity_temporal_metadata.entity_edition_id
                INNER JOIN ontology_ids ON ontology_ids.ontology_id = \
                 entity_is_of_type.entity_type_ontology_id
                WHERE ontology_ids.base_url = $1 AND ontology_ids.version = $2
                ORDER BY RANDOM()
                LIMIT 50
                ",
                &[&entity_type_id.base_url, &entity_type_id.version],
            )
            .instrument(tracing::info_span!(
                "SELECT",
                otel.kind = "client",
                db.system = "postgresql",
                peer.service = "Postgres",
            ))
            .await
            .unwrap_or_else(|err| {
                panic!("failed to sample entities for entity type `{entity_type_id}`: {err}");
            })
            .into_iter()