mz_persist_client/cli/
args.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
10//! CLI argument types for persist
11
12use crate::ShardId;
13use crate::cfg::PersistConfig;
14use crate::internal::metrics::{MetricsBlob, MetricsConsensus};
15use crate::internal::state_versions::StateVersions;
16use crate::metrics::Metrics;
17use async_trait::async_trait;
18use bytes::Bytes;
19use mz_build_info::BuildInfo;
20use mz_ore::bytes::SegmentedBytes;
21use mz_ore::metrics::MetricsRegistry;
22use mz_ore::now::SYSTEM_TIME;
23use mz_ore::url::SensitiveUrl;
24use mz_persist::cfg::{BlobConfig, ConsensusConfig};
25use mz_persist::location::{
26    Blob, BlobMetadata, CaSResult, Consensus, ExternalError, ResultStream, SeqNo, Tasked,
27    VersionedData,
28};
29use std::str::FromStr;
30use std::sync::Arc;
31use std::sync::atomic::{AtomicBool, Ordering};
32use tracing::warn;
33
34/// Arguments for commands that work over both backing stores.
35/// TODO: squish this into `StateArgs`.
36#[derive(Debug, Clone, clap::Parser)]
37pub struct StoreArgs {
38    /// Consensus to use.
39    ///
40    /// When connecting to a deployed environment's consensus table, the Postgres/CRDB connection
41    /// string must contain the database name and `options=--search_path=consensus`.
42    ///
43    /// When connecting to Cockroach Cloud, use the following format:
44    ///
45    /// ```text
46    /// postgresql://<user>:$COCKROACH_PW@<hostname>:<port>/environment_<environment-id>
47    ///   ?sslmode=verify-full
48    ///   &sslrootcert=/path/to/cockroach-cloud/certs/cluster-ca.crt
49    ///   &options=--search_path=consensus
50    /// ```
51    ///
52    #[clap(long, verbatim_doc_comment, env = "CONSENSUS_URI")]
53    pub(crate) consensus_uri: SensitiveUrl,
54
55    /// Blob to use
56    ///
57    /// When connecting to a deployed environment's blob, the necessary connection glue must be in
58    /// place. e.g. for S3, sign into SSO, set AWS_PROFILE and AWS_REGION appropriately, with a blob
59    /// URI scoped to the environment's bucket prefix.
60    #[clap(long, env = "BLOB_URI")]
61    pub(crate) blob_uri: SensitiveUrl,
62}
63
64/// Arguments for viewing the current state of a given shard
65#[derive(Debug, Clone, clap::Parser)]
66pub struct StateArgs {
67    /// Shard to view
68    #[clap(long)]
69    pub(crate) shard_id: String,
70
71    /// Consensus to use.
72    ///
73    /// When connecting to a deployed environment's consensus table, the Postgres/CRDB connection
74    /// string must contain the database name and `options=--search_path=consensus`.
75    ///
76    /// When connecting to Cockroach Cloud, use the following format:
77    ///
78    /// ```text
79    /// postgresql://<user>:$COCKROACH_PW@<hostname>:<port>/environment_<environment-id>
80    ///   ?sslmode=verify-full
81    ///   &sslrootcert=/path/to/cockroach-cloud/certs/cluster-ca.crt
82    ///   &options=--search_path=consensus
83    /// ```
84    ///
85    #[clap(long, verbatim_doc_comment, env = "CONSENSUS_URI")]
86    pub(crate) consensus_uri: SensitiveUrl,
87
88    /// Blob to use
89    ///
90    /// When connecting to a deployed environment's blob, the necessary connection glue must be in
91    /// place. e.g. for S3, sign into SSO, set AWS_PROFILE and AWS_REGION appropriately, with a blob
92    /// URI scoped to the environment's bucket prefix.
93    #[clap(long, env = "BLOB_URI")]
94    pub(crate) blob_uri: SensitiveUrl,
95}
96
97// BuildInfo with a larger version than any version we expect to see in prod,
98// to ensure that any data read is from a smaller version and does not trigger
99// alerts.
100pub(crate) const READ_ALL_BUILD_INFO: BuildInfo = BuildInfo {
101    version: "99.999.99+test",
102    sha: "0000000000000000000000000000000000000000",
103};
104
105// All `inspect` command are read-only.
106pub(crate) const NO_COMMIT: bool = false;
107
108impl StateArgs {
109    pub(crate) fn shard_id(&self) -> ShardId {
110        ShardId::from_str(&self.shard_id).expect("invalid shard id")
111    }
112
113    pub(crate) async fn open(&self) -> Result<StateVersions, anyhow::Error> {
114        let cfg = PersistConfig::new_default_configs(&READ_ALL_BUILD_INFO, SYSTEM_TIME.clone());
115        let metrics = Arc::new(Metrics::new(&cfg, &MetricsRegistry::new()));
116        let consensus =
117            make_consensus(&cfg, &self.consensus_uri, NO_COMMIT, Arc::clone(&metrics)).await?;
118        let blob = make_blob(&cfg, &self.blob_uri, NO_COMMIT, Arc::clone(&metrics)).await?;
119        Ok(StateVersions::new(cfg, consensus, blob, metrics))
120    }
121}
122
123pub(super) async fn make_consensus(
124    cfg: &PersistConfig,
125    consensus_uri: &SensitiveUrl,
126    commit: bool,
127    metrics: Arc<Metrics>,
128) -> anyhow::Result<Arc<dyn Consensus>> {
129    let consensus = ConsensusConfig::try_from(
130        consensus_uri,
131        Box::new(cfg.clone()),
132        metrics.postgres_consensus.clone(),
133        Arc::clone(&cfg.configs),
134    )?;
135    let consensus = consensus.clone().open().await?;
136    let consensus = if commit {
137        consensus
138    } else {
139        Arc::new(ReadOnly::new(consensus))
140    };
141    let consensus = Arc::new(MetricsConsensus::new(consensus, Arc::clone(&metrics)));
142    let consensus = Arc::new(Tasked(consensus));
143    Ok(consensus)
144}
145
146pub(super) async fn make_blob(
147    cfg: &PersistConfig,
148    blob_uri: &SensitiveUrl,
149    commit: bool,
150    metrics: Arc<Metrics>,
151) -> anyhow::Result<Arc<dyn Blob>> {
152    let blob =
153        BlobConfig::try_from(blob_uri, Box::new(cfg.clone()), metrics.s3_blob.clone()).await?;
154    let blob = blob.clone().open().await?;
155    let blob = if commit {
156        blob
157    } else {
158        Arc::new(ReadOnly::new(blob))
159    };
160    let blob = Arc::new(MetricsBlob::new(blob, Arc::clone(&metrics)));
161    let blob = Arc::new(Tasked(blob));
162    Ok(blob)
163}
164
165/// Wrap a lower-level service (Blob or Consensus) to make it read only.
166/// This is probably not elaborate enough to work in general -- folks may expect to read
167/// their own writes, among other things -- but it should handle the case of GC, where
168/// all reads finish before the writes begin.
169#[derive(Debug)]
170struct ReadOnly<T> {
171    store: T,
172    ignored_write: AtomicBool,
173}
174
175impl<T> ReadOnly<T> {
176    fn new(store: T) -> Self {
177        Self {
178            store,
179            ignored_write: AtomicBool::new(false),
180        }
181    }
182
183    fn ignored_write(&self) -> bool {
184        self.ignored_write.load(Ordering::SeqCst)
185    }
186
187    fn ignoring_write(&self) {
188        self.ignored_write.store(true, Ordering::SeqCst)
189    }
190}
191
192#[async_trait]
193impl Blob for ReadOnly<Arc<dyn Blob>> {
194    async fn get(&self, key: &str) -> Result<Option<SegmentedBytes>, ExternalError> {
195        if self.ignored_write() {
196            warn!("potentially-invalid get({key}) after ignored write");
197        }
198        self.store.get(key).await
199    }
200
201    async fn list_keys_and_metadata(
202        &self,
203        key_prefix: &str,
204        f: &mut (dyn FnMut(BlobMetadata) + Send + Sync),
205    ) -> Result<(), ExternalError> {
206        if self.ignored_write() {
207            warn!("potentially-invalid list_keys_and_metadata() after ignored write");
208        }
209        self.store.list_keys_and_metadata(key_prefix, f).await
210    }
211
212    async fn set(&self, key: &str, _value: Bytes) -> Result<(), ExternalError> {
213        warn!("ignoring set({key}) in read-only mode");
214        self.ignoring_write();
215        Ok(())
216    }
217
218    async fn delete(&self, key: &str) -> Result<Option<usize>, ExternalError> {
219        warn!("ignoring delete({key}) in read-only mode");
220        self.ignoring_write();
221        Ok(None)
222    }
223
224    async fn restore(&self, key: &str) -> Result<(), ExternalError> {
225        warn!("ignoring restore({key}) in read-only mode");
226        self.ignoring_write();
227        Ok(())
228    }
229}
230
231#[async_trait]
232impl Consensus for ReadOnly<Arc<dyn Consensus>> {
233    fn list_keys(&self) -> ResultStream<'_, String> {
234        if self.ignored_write() {
235            warn!("potentially-invalid list_keys() after ignored write");
236        }
237        self.store.list_keys()
238    }
239
240    async fn head(&self, key: &str) -> Result<Option<VersionedData>, ExternalError> {
241        if self.ignored_write() {
242            warn!("potentially-invalid head({key}) after ignored write");
243        }
244        self.store.head(key).await
245    }
246
247    async fn compare_and_set(
248        &self,
249        key: &str,
250        new: VersionedData,
251    ) -> Result<CaSResult, ExternalError> {
252        warn!(
253            "ignoring cas({key}) in read-only mode ({} bytes at seqno {:?})",
254            new.data.len(),
255            new.seqno,
256        );
257        self.ignoring_write();
258        Ok(CaSResult::Committed)
259    }
260
261    async fn scan(
262        &self,
263        key: &str,
264        from: SeqNo,
265        limit: usize,
266    ) -> Result<Vec<VersionedData>, ExternalError> {
267        if self.ignored_write() {
268            warn!("potentially-invalid scan({key}) after ignored write");
269        }
270        self.store.scan(key, from, limit).await
271    }
272
273    async fn truncate(&self, key: &str, seqno: SeqNo) -> Result<Option<usize>, ExternalError> {
274        warn!("ignoring truncate({key}) in read-only mode (to seqno {seqno})");
275        self.ignoring_write();
276        Ok(None)
277    }
278}
mz_persist_client/cli/args.rs

mz_persist_client/cli/
args.rs
