Skip to main content

mz_sql/session/
hint.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use internal::Private;
11
12/// Prometheus label for [`ApplicationNameHint::Unspecified`].
13const UNSPECIFIED_LABEL: &str = "unspecified";
14/// Prometheus label for [`ApplicationNameHint::Unrecognized`].
15const UNRECOGNIZED_LABEL: &str = "unrecognized";
16/// Prometheus label for [`ApplicationNameHint::Psql`].
17const PSQL_LABEL: &str = "psql";
18/// Prometheus label for [`ApplicationNameHint::Dbt`].
19const DBT_LABEL: &str = "dbt";
20/// Prometheus label for [`ApplicationNameHint::WebConsole`].
21const WEB_CONSOLE_LABEL: &str = "web_console";
22/// Promehteus label for [`ApplicationNameHint::WebConsoleShell`].
23const WEB_CONSOLE_SHELL_LABEL: &str = "web_console_shell";
24/// Prometheus label for [`ApplicationNameHint::MzPsql`].
25const MZ_PSQL_LABEL: &str = "mz_psql";
26/// Prometheus label for [`ApplicationNameHint::MaterializeFivetranDestination`].
27const MATERIALIZE_FIVETRAN_DESTINATION_LABEL: &str = "materialize_fivetran_destination";
28/// Prometheus label for [`ApplicationNameHint::TerraformProviderMaterialize`].
29const TERRAFORM_PROVIDER_MATERIALIZE_LABEL: &str = "terraform_provider_materialize";
30/// Prometheus label for [`ApplicationNameHint::TablePlus`].
31const TABLE_PLUS_LABEL: &str = "table_plus";
32/// Prometheus label for [`ApplicationNameHint::DataGrip`].
33const DATA_GRIP_LABEL: &str = "data_grip";
34/// Prometheus label for [`ApplicationNameHint::DBeaver`].
35const D_BEAVER_LABEL: &str = "dbeaver";
36/// Prometheus label for [`ApplicationNameHint::MzVscode`].
37const MZ_VSCODE_LABEL: &str = "mz_vscode";
38/// Prometheus label for [`ApplicationNameHint::MzGrafanaIntegration`].
39const MZ_GRAFANA_LABEL: &str = "mz_grafana";
40/// Prometheus label for [`ApplicationNameHint::MzMcpAgent`].
41const MZ_MCP_AGENT_LABEL: &str = "mz_mcp_agents";
42/// Prometheus label for [`ApplicationNameHint::MzMcpDeveloper`].
43const MZ_MCP_DEVELOPER_LABEL: &str = "mz_mcp_developer";
44
45/// A hint for what application is making a request to the adapter.
46///
47/// Note: [`ApplicationNameHint`] gets logged as a label in our Prometheus metrics, and for
48/// labels we need to be conscious of the cardinality, so please be careful with how many
49/// variants we add to this enum.
50///
51/// Note: each enum variant contains an `internal::Private` to prevent creating this enum
52/// directly. To create an instance of [`ApplicationNameHint`] please see
53/// [`ApplicationNameHint::from_str`].
54#[derive(Debug, Copy, Clone)]
55pub enum ApplicationNameHint {
56    /// No `application_name` was set.
57    Unspecified(Private),
58    /// An `application_name` was set, but it's not one we recognize.
59    Unrecognized(Private),
60    /// Request came from `psql`.
61    Psql(Private),
62    /// Request came from `dbt`.
63    Dbt(Private),
64    /// Request came from our web console.
65    WebConsole(Private),
66    /// Request came from the SQL shell in our web console.
67    WebConsoleShell(Private),
68    /// Request came from the `psql` shell spawned by `mz`.
69    MzPsql(Private),
70    /// Request came from our Fivetran Destination,
71    MaterializeFivetranDestination(Private),
72    /// Request came from a version of our Terraform provider.
73    TerraformProviderMaterialize(Private),
74    /// Request came from TablePlus.
75    TablePlus(Private),
76    /// Request came from a version of DataGrip.
77    DataGrip(Private),
78    /// Request came from a version of DBeaver.
79    DBeaver(Private),
80    /// Request came from our Visual Studio Code integration.
81    MzVscode(Private),
82    /// Request came from our Grafana integration.
83    MzGrafanaIntegration(Private),
84    /// Request came from the MCP agent endpoint.
85    MzMcpAgent(Private),
86    /// Request came from the MCP developer endpoint.
87    MzMcpDeveloper(Private),
88}
89
90impl ApplicationNameHint {
91    pub fn from_str(s: &str) -> Self {
92        match s.to_lowercase().as_str() {
93            "psql" => ApplicationNameHint::Psql(Private),
94            "dbt" => ApplicationNameHint::Dbt(Private),
95            "web_console" => ApplicationNameHint::WebConsole(Private),
96            "web_console_shell" => ApplicationNameHint::WebConsoleShell(Private),
97            "mz_psql" => ApplicationNameHint::MzPsql(Private),
98            // Note: Make sure this is kept in sync with the `fivetran-destination` crate.
99            "materialize_fivetran_destination" => {
100                ApplicationNameHint::MaterializeFivetranDestination(Private)
101            }
102            "tableplus" => ApplicationNameHint::TablePlus(Private),
103            "mz_vscode" => ApplicationNameHint::MzVscode(Private),
104            "mz_grafana_integration" => ApplicationNameHint::MzGrafanaIntegration(Private),
105            "mz_mcp_agents" => ApplicationNameHint::MzMcpAgent(Private),
106            "mz_mcp_developer" => ApplicationNameHint::MzMcpDeveloper(Private),
107            // Terraform provides the version as a suffix.
108            x if x.starts_with("terraform-provider-materialize") => {
109                ApplicationNameHint::TerraformProviderMaterialize(Private)
110            }
111            // DataGrip provides the version as a suffix.
112            x if x.starts_with("datagrip") => ApplicationNameHint::DataGrip(Private),
113            // DBeaver provides the version as a suffix.
114            x if x.starts_with("dbeaver") => ApplicationNameHint::DBeaver(Private),
115            "" => ApplicationNameHint::Unspecified(Private),
116            // TODO(parkertimmerman): We should keep some record of these "unrecognized"
117            // names, and possibly support more popular ones in the future.
118            _ => ApplicationNameHint::Unrecognized(Private),
119        }
120    }
121
122    pub fn as_str(&self) -> &'static str {
123        match self {
124            ApplicationNameHint::Unspecified(_) => UNSPECIFIED_LABEL,
125            ApplicationNameHint::Unrecognized(_) => UNRECOGNIZED_LABEL,
126            ApplicationNameHint::Psql(_) => PSQL_LABEL,
127            ApplicationNameHint::Dbt(_) => DBT_LABEL,
128            ApplicationNameHint::WebConsole(_) => WEB_CONSOLE_LABEL,
129            ApplicationNameHint::WebConsoleShell(_) => WEB_CONSOLE_SHELL_LABEL,
130            ApplicationNameHint::MzPsql(_) => MZ_PSQL_LABEL,
131            ApplicationNameHint::MaterializeFivetranDestination(_) => {
132                MATERIALIZE_FIVETRAN_DESTINATION_LABEL
133            }
134            ApplicationNameHint::TablePlus(_) => TABLE_PLUS_LABEL,
135            ApplicationNameHint::MzVscode(_) => MZ_VSCODE_LABEL,
136            ApplicationNameHint::MzGrafanaIntegration(_) => MZ_GRAFANA_LABEL,
137            ApplicationNameHint::MzMcpAgent(_) => MZ_MCP_AGENT_LABEL,
138            ApplicationNameHint::MzMcpDeveloper(_) => MZ_MCP_DEVELOPER_LABEL,
139            ApplicationNameHint::TerraformProviderMaterialize(_) => {
140                TERRAFORM_PROVIDER_MATERIALIZE_LABEL
141            }
142            ApplicationNameHint::DataGrip(_) => DATA_GRIP_LABEL,
143            ApplicationNameHint::DBeaver(_) => D_BEAVER_LABEL,
144        }
145    }
146
147    /// Returns whether or not we should trace errors for this requests with this application name.
148    pub fn should_trace_errors(&self) -> bool {
149        // Note(parkmycar): For now we only trace errors for the web console since we contol all of
150        // those queries and in general they should never fail. As opposed to user queries which
151        // are arbitrary.
152        matches!(self, ApplicationNameHint::WebConsole(_))
153    }
154}
155
156mod internal {
157    #[derive(Debug, Copy, Clone)]
158    pub struct Private;
159}