scouter_sql/sql/
query.rs

1//constants
2
3// spc
4const INSERT_DRIFT_RECORD: &str = include_str!("scripts/spc/insert_spc_drift_record.sql");
5const INSERT_SPC_DRIFT_RECORD_BATCH: &str =
6    include_str!("scripts/spc/insert_spc_drift_record_batch.sql");
7const GET_SPC_FEATURES: &str = include_str!("scripts/spc/unique_spc_features.sql");
8const GET_BINNED_SPC_FEATURE_VALUES: &str =
9    include_str!("scripts/spc/binned_spc_feature_values.sql");
10const GET_SPC_FEATURE_VALUES: &str = include_str!("scripts/spc/get_spc_feature_values.sql");
11const GET_SPC_ENTITIES: &str = include_str!("scripts/spc/get_spc_entities_for_archive.sql");
12const GET_SPC_DATA_FOR_ARCHIVE: &str = include_str!("scripts/spc/get_spc_data_for_archive.sql");
13const UPDATE_SPC_ENTITIES: &str = include_str!("scripts/spc/update_data_to_archived.sql");
14
15// psi
16const INSERT_BIN_COUNTS: &str = include_str!("scripts/psi/insert_bin_counts.sql");
17const INSERT_BIN_COUNTS_BATCH: &str = include_str!("scripts/psi/insert_bin_counts_batch.sql");
18const GET_BINNED_PSI_FEATURE_BINS: &str =
19    include_str!("scripts/psi/binned_psi_feature_bin_proportions.sql");
20const GET_FEATURE_BIN_PROPORTIONS: &str =
21    include_str!("scripts/psi/get_feature_bin_proportions.sql");
22const GET_BIN_COUNT_ENTITIES: &str =
23    include_str!("scripts/psi/get_bin_count_entities_for_archive.sql");
24const GET_BIN_COUNT_DATA_FOR_ARCHIVE: &str =
25    include_str!("scripts/psi/get_bin_count_data_for_archive.sql");
26const UPDATE_BIN_COUNT_ENTITIES: &str = include_str!("scripts/psi/update_data_to_archived.sql");
27
28// custom
29const GET_BINNED_CUSTOM_METRIC_VALUES: &str =
30    include_str!("scripts/custom/binned_custom_metric_values.sql");
31const GET_CUSTOM_METRIC_VALUES: &str = include_str!("scripts/custom/get_custom_metric_values.sql");
32const INSERT_CUSTOM_METRIC_VALUES: &str =
33    include_str!("scripts/custom/insert_custom_metric_values.sql");
34const INSERT_CUSTOM_METRIC_VALUES_BATCH: &str =
35    include_str!("scripts/custom/insert_custom_metric_values_batch.sql");
36const GET_CUSTOM_ENTITIES: &str =
37    include_str!("scripts/custom/get_custom_metric_entities_for_archive.sql");
38const GET_CUSTOM_DATA_FOR_ARCHIVE: &str =
39    include_str!("scripts/custom/get_custom_metric_data_for_archive.sql");
40const UPDATE_CUSTOM_ENTITIES: &str = include_str!("scripts/custom/update_data_to_archived.sql");
41
42// llm
43const GET_LLM_METRIC_VALUES: &str = include_str!("scripts/llm/get_llm_metric_values.sql");
44const GET_BINNED_LLM_METRIC_VALUES: &str = include_str!("scripts/llm/binned_llm_metric_values.sql");
45const INSERT_LLM_METRIC_VALUES_BATCH: &str =
46    include_str!("scripts/llm/insert_llm_metric_values.sql");
47const INSERT_LLM_DRIFT_RECORD: &str = include_str!("scripts/llm/insert_llm_drift_record.sql");
48
49const GET_LLM_DRIFT_RECORDS: &str = include_str!("scripts/llm/get_llm_drift_records.sql");
50const UPDATE_LLM_DRIFT_TASK: &str = include_str!("scripts/llm/update_llm_drift_record.sql");
51const GET_LLM_DRIFT_RECORD_ENTITIES: &str =
52    include_str!("scripts/llm/get_llm_drift_record_entities_for_archive.sql");
53const GET_LLM_METRIC_ENTITIES: &str =
54    include_str!("scripts/llm/get_llm_metric_entities_for_archive.sql");
55const GET_LLM_DRIFT_RECORD_DATA_FOR_ARCHIVE: &str =
56    include_str!("scripts/llm/get_llm_drift_record_data_for_archive.sql");
57const GET_LLM_METRIC_DATA_FOR_ARCHIVE: &str =
58    include_str!("scripts/llm/get_llm_metric_data_for_archive.sql");
59const UPDATE_LLM_METRIC_ENTITIES: &str =
60    include_str!("scripts/llm/update_llm_drift_metric_to_archived.sql");
61const UPDATE_LLM_DRIFT_ENTITIES: &str =
62    include_str!("scripts/llm/update_llm_drift_record_to_archived.sql");
63
64// observability (experimental)
65const GET_BINNED_OBSERVABILITY_METRICS: &str =
66    include_str!("scripts/observability/binned_observability_metrics.sql");
67const INSERT_OBSERVABILITY_RECORD: &str =
68    include_str!("scripts/observability/insert_observability_record.sql");
69
70//profile
71const INSERT_DRIFT_PROFILE: &str = include_str!("scripts/profile/insert_drift_profile.sql");
72const GET_DRIFT_PROFILE: &str = include_str!("scripts/profile/get_drift_profile.sql");
73const UPDATE_DRIFT_PROFILE_RUN_DATES: &str =
74    include_str!("scripts/profile/update_drift_profile_run_dates.sql");
75const UPDATE_DRIFT_PROFILE_STATUS: &str =
76    include_str!("scripts/profile/update_drift_profile_status.sql");
77const UPDATE_DRIFT_PROFILE: &str = include_str!("scripts/profile/update_drift_profile.sql");
78const DEACTIVATE_DRIFT_PROFILES: &str =
79    include_str!("scripts/profile/deactivate_drift_profiles.sql");
80const GET_PROFILE_VERSIONS: &str = include_str!("scripts/profile/get_profile_versions.sql");
81
82// alert
83const INSERT_DRIFT_ALERT: &str = include_str!("scripts/alert/insert_drift_alert.sql");
84const GET_DRIFT_ALERTS: &str = include_str!("scripts/alert/get_drift_alerts.sql");
85const UPDATE_ALERT_STATUS: &str = include_str!("scripts/alert/update_alert_status.sql");
86
87// poll
88const GET_DRIFT_TASK: &str = include_str!("scripts/poll/poll_for_drift_task.sql");
89const GET_PENDING_LLM_DRIFT_TASK: &str = include_str!("scripts/poll/poll_for_llm_drift_task.sql");
90
91// auth
92const INSERT_USER: &str = include_str!("scripts/user/insert_user.sql");
93const GET_USER: &str = include_str!("scripts/user/get_user.sql");
94const UPDATE_USER: &str = include_str!("scripts/user/update_user.sql");
95const GET_USERS: &str = include_str!("scripts/user/get_users.sql");
96const LAST_ADMIN: &str = include_str!("scripts/user/last_admin.sql");
97const DELETE_USER: &str = include_str!("scripts/user/delete_user.sql");
98
99#[allow(dead_code)]
100pub enum Queries {
101    GetSpcFeatures,
102    InsertDriftRecord,
103    InsertBinCounts,
104    InsertDriftProfile,
105    InsertDriftAlert,
106    InsertObservabilityRecord,
107    GetDriftAlerts,
108    GetBinnedSpcFeatureValues,
109    GetBinnedPsiFeatureBins,
110    GetBinnedMetricValues,
111    GetBinnedObservabilityMetrics,
112    GetSpcFeatureValues,
113    GetDriftTask,
114    GetDriftProfile,
115    UpdateDriftProfileRunDates,
116    UpdateDriftProfileStatus,
117    DeactivateDriftProfiles,
118    UpdateDriftProfile,
119    GetFeatureBinProportions,
120    GetCustomMetricValues,
121    InsertCustomMetricValues,
122
123    InsertCustomMetricValuesBatch,
124    InsertSpcDriftRecordBatch,
125    InsertBinCountsBatch,
126
127    // archive
128    // entities
129    GetBinCountEntities,
130    GetCustomEntities,
131    GetSpcEntities,
132
133    // data
134    GetBinCountDataForArchive,
135    GetCustomDataForArchive,
136    GetSpcDataForArchive,
137
138    // update
139    UpdateBinCountEntities,
140    UpdateCustomEntities,
141    UpdateSpcEntities,
142
143    // user
144    InsertUser,
145    GetUser,
146    UpdateUser,
147    GetUsers,
148    LastAdmin,
149    DeleteUser,
150    UpdateAlertStatus,
151
152    // llm
153    GetLLMMetricValues,
154    GetLLMDriftRecords,
155    GetBinnedMetrics,
156    InsertLLMMetricValuesBatch,
157    InsertLLMDriftRecord,
158    GetPendingLLMDriftTask,
159    UpdateLLMDriftTask,
160
161    GetLLMDriftRecordEntitiesForArchive,
162    GetLLMMetricEntitiesForArchive,
163    GetLLMDriftRecordDataForArchive,
164    GetLLMMetricDataForArchive,
165
166    UpdateLLMMetricEntities,
167    UpdateLLMDriftEntities,
168
169    // profile
170    GetProfileVersions,
171}
172
173impl Queries {
174    // TODO: shouldn't we just return the string directly? Not sure if that's true for all db operations, I'm
175    // just noticing it in the few that im working on. (user related queries)
176    pub fn get_query(&self) -> SqlQuery {
177        match self {
178            // load sql file from scripts/insert.sql
179            Queries::GetSpcFeatures => SqlQuery::new(GET_SPC_FEATURES),
180            Queries::InsertDriftRecord => SqlQuery::new(INSERT_DRIFT_RECORD),
181            Queries::GetBinnedSpcFeatureValues => SqlQuery::new(GET_BINNED_SPC_FEATURE_VALUES),
182            Queries::GetBinnedPsiFeatureBins => SqlQuery::new(GET_BINNED_PSI_FEATURE_BINS),
183            Queries::GetBinnedMetricValues => SqlQuery::new(GET_BINNED_CUSTOM_METRIC_VALUES),
184            Queries::GetBinnedObservabilityMetrics => {
185                SqlQuery::new(GET_BINNED_OBSERVABILITY_METRICS)
186            }
187            Queries::GetSpcFeatureValues => SqlQuery::new(GET_SPC_FEATURE_VALUES),
188            Queries::InsertDriftProfile => SqlQuery::new(INSERT_DRIFT_PROFILE),
189            Queries::InsertDriftAlert => SqlQuery::new(INSERT_DRIFT_ALERT),
190            Queries::InsertObservabilityRecord => SqlQuery::new(INSERT_OBSERVABILITY_RECORD),
191            Queries::GetDriftAlerts => SqlQuery::new(GET_DRIFT_ALERTS),
192            Queries::GetDriftTask => SqlQuery::new(GET_DRIFT_TASK),
193            Queries::UpdateDriftProfileRunDates => SqlQuery::new(UPDATE_DRIFT_PROFILE_RUN_DATES),
194            Queries::UpdateDriftProfileStatus => SqlQuery::new(UPDATE_DRIFT_PROFILE_STATUS),
195            Queries::UpdateDriftProfile => SqlQuery::new(UPDATE_DRIFT_PROFILE),
196            Queries::DeactivateDriftProfiles => SqlQuery::new(DEACTIVATE_DRIFT_PROFILES),
197            Queries::GetDriftProfile => SqlQuery::new(GET_DRIFT_PROFILE),
198            Queries::GetFeatureBinProportions => SqlQuery::new(GET_FEATURE_BIN_PROPORTIONS),
199            Queries::InsertBinCounts => SqlQuery::new(INSERT_BIN_COUNTS),
200            Queries::GetCustomMetricValues => SqlQuery::new(GET_CUSTOM_METRIC_VALUES),
201            Queries::InsertCustomMetricValues => SqlQuery::new(INSERT_CUSTOM_METRIC_VALUES),
202            Queries::GetBinCountEntities => SqlQuery::new(GET_BIN_COUNT_ENTITIES),
203            Queries::GetCustomEntities => SqlQuery::new(GET_CUSTOM_ENTITIES),
204            Queries::GetSpcEntities => SqlQuery::new(GET_SPC_ENTITIES),
205            Queries::GetBinCountDataForArchive => SqlQuery::new(GET_BIN_COUNT_DATA_FOR_ARCHIVE),
206            Queries::GetCustomDataForArchive => SqlQuery::new(GET_CUSTOM_DATA_FOR_ARCHIVE),
207            Queries::GetSpcDataForArchive => SqlQuery::new(GET_SPC_DATA_FOR_ARCHIVE),
208            Queries::UpdateBinCountEntities => SqlQuery::new(UPDATE_BIN_COUNT_ENTITIES),
209            Queries::UpdateCustomEntities => SqlQuery::new(UPDATE_CUSTOM_ENTITIES),
210            Queries::UpdateSpcEntities => SqlQuery::new(UPDATE_SPC_ENTITIES),
211            Queries::GetProfileVersions => SqlQuery::new(GET_PROFILE_VERSIONS),
212
213            Queries::InsertUser => SqlQuery::new(INSERT_USER),
214            Queries::GetUser => SqlQuery::new(GET_USER),
215            Queries::UpdateUser => SqlQuery::new(UPDATE_USER),
216            Queries::GetUsers => SqlQuery::new(GET_USERS),
217            Queries::LastAdmin => SqlQuery::new(LAST_ADMIN),
218            Queries::DeleteUser => SqlQuery::new(DELETE_USER),
219            Queries::UpdateAlertStatus => SqlQuery::new(UPDATE_ALERT_STATUS),
220
221            //llm
222            Queries::GetLLMMetricValues => SqlQuery::new(GET_LLM_METRIC_VALUES),
223            Queries::GetBinnedMetrics => SqlQuery::new(GET_BINNED_LLM_METRIC_VALUES),
224            Queries::InsertLLMMetricValuesBatch => SqlQuery::new(INSERT_LLM_METRIC_VALUES_BATCH),
225            Queries::InsertLLMDriftRecord => SqlQuery::new(INSERT_LLM_DRIFT_RECORD),
226
227            Queries::GetLLMDriftRecords => SqlQuery::new(GET_LLM_DRIFT_RECORDS),
228            Queries::GetPendingLLMDriftTask => SqlQuery::new(GET_PENDING_LLM_DRIFT_TASK),
229            Queries::GetLLMDriftRecordEntitiesForArchive => {
230                SqlQuery::new(GET_LLM_DRIFT_RECORD_ENTITIES)
231            }
232            Queries::GetLLMMetricEntitiesForArchive => SqlQuery::new(GET_LLM_METRIC_ENTITIES),
233            Queries::GetLLMDriftRecordDataForArchive => {
234                SqlQuery::new(GET_LLM_DRIFT_RECORD_DATA_FOR_ARCHIVE)
235            }
236            Queries::GetLLMMetricDataForArchive => SqlQuery::new(GET_LLM_METRIC_DATA_FOR_ARCHIVE),
237            Queries::UpdateLLMMetricEntities => SqlQuery::new(UPDATE_LLM_METRIC_ENTITIES),
238            Queries::UpdateLLMDriftEntities => SqlQuery::new(UPDATE_LLM_DRIFT_ENTITIES),
239
240            Queries::InsertCustomMetricValuesBatch => {
241                SqlQuery::new(INSERT_CUSTOM_METRIC_VALUES_BATCH)
242            }
243            Queries::InsertSpcDriftRecordBatch => SqlQuery::new(INSERT_SPC_DRIFT_RECORD_BATCH),
244            Queries::InsertBinCountsBatch => SqlQuery::new(INSERT_BIN_COUNTS_BATCH),
245            Queries::UpdateLLMDriftTask => SqlQuery::new(UPDATE_LLM_DRIFT_TASK),
246        }
247    }
248}
249
250pub struct SqlQuery {
251    pub sql: String,
252}
253
254impl SqlQuery {
255    fn new(sql: &str) -> Self {
256        Self {
257            sql: sql.to_string(),
258        }
259    }
260}
scouter_sql/sql/query.rs

scouter_sql/sql/
query.rs
