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");
81const LIST_DRIFT_PROFILES: &str = include_str!("scripts/profile/list_drift_profiles.sql");
82
83// alert
84const INSERT_DRIFT_ALERT: &str = include_str!("scripts/alert/insert_drift_alert.sql");
85const GET_DRIFT_ALERTS: &str = include_str!("scripts/alert/get_drift_alerts.sql");
86const UPDATE_ALERT_STATUS: &str = include_str!("scripts/alert/update_alert_status.sql");
87
88// poll
89const GET_DRIFT_TASK: &str = include_str!("scripts/poll/poll_for_drift_task.sql");
90const GET_PENDING_LLM_DRIFT_TASK: &str = include_str!("scripts/poll/poll_for_llm_drift_task.sql");
91
92// auth
93const INSERT_USER: &str = include_str!("scripts/user/insert_user.sql");
94const GET_USER: &str = include_str!("scripts/user/get_user.sql");
95const UPDATE_USER: &str = include_str!("scripts/user/update_user.sql");
96const GET_USERS: &str = include_str!("scripts/user/get_users.sql");
97const LAST_ADMIN: &str = include_str!("scripts/user/last_admin.sql");
98const DELETE_USER: &str = include_str!("scripts/user/delete_user.sql");
99
100// trace
101const UPSERT_TRACE: &str = include_str!("scripts/trace/upsert_trace.sql");
102const INSERT_TRACE_SPAN: &str = include_str!("scripts/trace/insert_span.sql");
103const INSERT_TRACE_BAGGAGE: &str = include_str!("scripts/trace/insert_baggage.sql");
104const GET_PAGINATED_TRACES: &str = include_str!("scripts/trace/get_paginated_traces.sql");
105const GET_TRACE_SPANS: &str = include_str!("scripts/trace/get_trace_spans.sql");
106const GET_TRACE_METRICS: &str = include_str!("scripts/trace/get_trace_metrics.sql");
107const GET_TRACE_BAGGAGE: &str = include_str!("scripts/trace/get_trace_baggage.sql");
108
109// tags
110const INSERT_TAG: &str = include_str!("scripts/tag/insert_tags.sql");
111const GET_TAGS: &str = include_str!("scripts/tag/get_tags.sql");
112
113#[allow(dead_code)]
114pub enum Queries {
115    GetSpcFeatures,
116    InsertDriftRecord,
117    InsertBinCounts,
118    InsertDriftProfile,
119    InsertDriftAlert,
120    InsertObservabilityRecord,
121    GetDriftAlerts,
122    GetBinnedSpcFeatureValues,
123    GetBinnedPsiFeatureBins,
124    GetBinnedMetricValues,
125    GetBinnedObservabilityMetrics,
126    GetSpcFeatureValues,
127    GetDriftTask,
128    GetDriftProfile,
129    UpdateDriftProfileRunDates,
130    UpdateDriftProfileStatus,
131    DeactivateDriftProfiles,
132    UpdateDriftProfile,
133    GetFeatureBinProportions,
134    GetCustomMetricValues,
135    InsertCustomMetricValues,
136
137    InsertCustomMetricValuesBatch,
138    InsertSpcDriftRecordBatch,
139    InsertBinCountsBatch,
140
141    // archive
142    // entities
143    GetBinCountEntities,
144    GetCustomEntities,
145    GetSpcEntities,
146
147    // data
148    GetBinCountDataForArchive,
149    GetCustomDataForArchive,
150    GetSpcDataForArchive,
151
152    // update
153    UpdateBinCountEntities,
154    UpdateCustomEntities,
155    UpdateSpcEntities,
156
157    // user
158    InsertUser,
159    GetUser,
160    UpdateUser,
161    GetUsers,
162    LastAdmin,
163    DeleteUser,
164    UpdateAlertStatus,
165
166    // llm
167    GetLLMMetricValues,
168    GetLLMDriftRecords,
169    GetBinnedMetrics,
170    InsertLLMMetricValuesBatch,
171    InsertLLMDriftRecord,
172    GetPendingLLMDriftTask,
173    UpdateLLMDriftTask,
174
175    GetLLMDriftRecordEntitiesForArchive,
176    GetLLMMetricEntitiesForArchive,
177    GetLLMDriftRecordDataForArchive,
178    GetLLMMetricDataForArchive,
179
180    UpdateLLMMetricEntities,
181    UpdateLLMDriftEntities,
182
183    // profile
184    GetProfileVersions,
185    ListDriftProfiles,
186
187    //trace
188    UpsertTrace,
189    InsertTraceSpan,
190    InsertTraceBaggage,
191    GetPaginatedTraces,
192    GetTraceSpans,
193    GetTraceMetrics,
194    GetTraceBaggage,
195
196    // tags
197    InsertTag,
198    GetTags,
199}
200
201impl Queries {
202    // TODO: shouldn't we just return the string directly? Not sure if that's true for all db operations, I'm
203    // just noticing it in the few that im working on. (user related queries)
204    pub fn get_query(&self) -> SqlQuery {
205        match self {
206            // load sql file from scripts/insert.sql
207            Queries::GetSpcFeatures => SqlQuery::new(GET_SPC_FEATURES),
208            Queries::InsertDriftRecord => SqlQuery::new(INSERT_DRIFT_RECORD),
209            Queries::GetBinnedSpcFeatureValues => SqlQuery::new(GET_BINNED_SPC_FEATURE_VALUES),
210            Queries::GetBinnedPsiFeatureBins => SqlQuery::new(GET_BINNED_PSI_FEATURE_BINS),
211            Queries::GetBinnedMetricValues => SqlQuery::new(GET_BINNED_CUSTOM_METRIC_VALUES),
212            Queries::GetBinnedObservabilityMetrics => {
213                SqlQuery::new(GET_BINNED_OBSERVABILITY_METRICS)
214            }
215            Queries::GetSpcFeatureValues => SqlQuery::new(GET_SPC_FEATURE_VALUES),
216            Queries::InsertDriftProfile => SqlQuery::new(INSERT_DRIFT_PROFILE),
217            Queries::InsertDriftAlert => SqlQuery::new(INSERT_DRIFT_ALERT),
218            Queries::InsertObservabilityRecord => SqlQuery::new(INSERT_OBSERVABILITY_RECORD),
219            Queries::GetDriftAlerts => SqlQuery::new(GET_DRIFT_ALERTS),
220            Queries::GetDriftTask => SqlQuery::new(GET_DRIFT_TASK),
221            Queries::UpdateDriftProfileRunDates => SqlQuery::new(UPDATE_DRIFT_PROFILE_RUN_DATES),
222            Queries::UpdateDriftProfileStatus => SqlQuery::new(UPDATE_DRIFT_PROFILE_STATUS),
223            Queries::UpdateDriftProfile => SqlQuery::new(UPDATE_DRIFT_PROFILE),
224            Queries::DeactivateDriftProfiles => SqlQuery::new(DEACTIVATE_DRIFT_PROFILES),
225            Queries::GetDriftProfile => SqlQuery::new(GET_DRIFT_PROFILE),
226            Queries::GetFeatureBinProportions => SqlQuery::new(GET_FEATURE_BIN_PROPORTIONS),
227            Queries::InsertBinCounts => SqlQuery::new(INSERT_BIN_COUNTS),
228            Queries::GetCustomMetricValues => SqlQuery::new(GET_CUSTOM_METRIC_VALUES),
229            Queries::InsertCustomMetricValues => SqlQuery::new(INSERT_CUSTOM_METRIC_VALUES),
230            Queries::GetBinCountEntities => SqlQuery::new(GET_BIN_COUNT_ENTITIES),
231            Queries::GetCustomEntities => SqlQuery::new(GET_CUSTOM_ENTITIES),
232            Queries::GetSpcEntities => SqlQuery::new(GET_SPC_ENTITIES),
233            Queries::GetBinCountDataForArchive => SqlQuery::new(GET_BIN_COUNT_DATA_FOR_ARCHIVE),
234            Queries::GetCustomDataForArchive => SqlQuery::new(GET_CUSTOM_DATA_FOR_ARCHIVE),
235            Queries::GetSpcDataForArchive => SqlQuery::new(GET_SPC_DATA_FOR_ARCHIVE),
236            Queries::UpdateBinCountEntities => SqlQuery::new(UPDATE_BIN_COUNT_ENTITIES),
237            Queries::UpdateCustomEntities => SqlQuery::new(UPDATE_CUSTOM_ENTITIES),
238            Queries::UpdateSpcEntities => SqlQuery::new(UPDATE_SPC_ENTITIES),
239            Queries::GetProfileVersions => SqlQuery::new(GET_PROFILE_VERSIONS),
240            Queries::ListDriftProfiles => SqlQuery::new(LIST_DRIFT_PROFILES),
241
242            Queries::InsertUser => SqlQuery::new(INSERT_USER),
243            Queries::GetUser => SqlQuery::new(GET_USER),
244            Queries::UpdateUser => SqlQuery::new(UPDATE_USER),
245            Queries::GetUsers => SqlQuery::new(GET_USERS),
246            Queries::LastAdmin => SqlQuery::new(LAST_ADMIN),
247            Queries::DeleteUser => SqlQuery::new(DELETE_USER),
248            Queries::UpdateAlertStatus => SqlQuery::new(UPDATE_ALERT_STATUS),
249
250            //llm
251            Queries::GetLLMMetricValues => SqlQuery::new(GET_LLM_METRIC_VALUES),
252            Queries::GetBinnedMetrics => SqlQuery::new(GET_BINNED_LLM_METRIC_VALUES),
253            Queries::InsertLLMMetricValuesBatch => SqlQuery::new(INSERT_LLM_METRIC_VALUES_BATCH),
254            Queries::InsertLLMDriftRecord => SqlQuery::new(INSERT_LLM_DRIFT_RECORD),
255
256            Queries::GetLLMDriftRecords => SqlQuery::new(GET_LLM_DRIFT_RECORDS),
257            Queries::GetPendingLLMDriftTask => SqlQuery::new(GET_PENDING_LLM_DRIFT_TASK),
258            Queries::GetLLMDriftRecordEntitiesForArchive => {
259                SqlQuery::new(GET_LLM_DRIFT_RECORD_ENTITIES)
260            }
261            Queries::GetLLMMetricEntitiesForArchive => SqlQuery::new(GET_LLM_METRIC_ENTITIES),
262            Queries::GetLLMDriftRecordDataForArchive => {
263                SqlQuery::new(GET_LLM_DRIFT_RECORD_DATA_FOR_ARCHIVE)
264            }
265            Queries::GetLLMMetricDataForArchive => SqlQuery::new(GET_LLM_METRIC_DATA_FOR_ARCHIVE),
266            Queries::UpdateLLMMetricEntities => SqlQuery::new(UPDATE_LLM_METRIC_ENTITIES),
267            Queries::UpdateLLMDriftEntities => SqlQuery::new(UPDATE_LLM_DRIFT_ENTITIES),
268
269            Queries::InsertCustomMetricValuesBatch => {
270                SqlQuery::new(INSERT_CUSTOM_METRIC_VALUES_BATCH)
271            }
272            Queries::InsertSpcDriftRecordBatch => SqlQuery::new(INSERT_SPC_DRIFT_RECORD_BATCH),
273            Queries::InsertBinCountsBatch => SqlQuery::new(INSERT_BIN_COUNTS_BATCH),
274            Queries::UpdateLLMDriftTask => SqlQuery::new(UPDATE_LLM_DRIFT_TASK),
275            // trace
276            Queries::UpsertTrace => SqlQuery::new(UPSERT_TRACE),
277            Queries::InsertTraceSpan => SqlQuery::new(INSERT_TRACE_SPAN),
278            Queries::InsertTraceBaggage => SqlQuery::new(INSERT_TRACE_BAGGAGE),
279            Queries::GetPaginatedTraces => SqlQuery::new(GET_PAGINATED_TRACES),
280            Queries::GetTraceSpans => SqlQuery::new(GET_TRACE_SPANS),
281            Queries::GetTraceMetrics => SqlQuery::new(GET_TRACE_METRICS),
282            Queries::GetTraceBaggage => SqlQuery::new(GET_TRACE_BAGGAGE),
283            // tags
284            Queries::InsertTag => SqlQuery::new(INSERT_TAG),
285            Queries::GetTags => SqlQuery::new(GET_TAGS),
286        }
287    }
288}
289
290pub struct SqlQuery {
291    pub sql: String,
292}
293
294impl SqlQuery {
295    fn new(sql: &str) -> Self {
296        Self {
297            sql: sql.to_string(),
298        }
299    }
300}
scouter_sql/sql/query.rs

scouter_sql/sql/
query.rs
