dbcrossbarlib/drivers/bigquery_shared/
table_name.rs

1//! BigQuery table names.
2
3use lazy_static::lazy_static;
4use regex::Regex;
5use std::{fmt, str::FromStr};
6
7use crate::common::*;
8use crate::drivers::bigquery::BigQueryLocator;
9
10/// A BigQuery table name of the form `"project:dataset.table"`.
11#[derive(Clone, Debug, Eq, PartialEq)]
12pub(crate) struct TableName {
13    /// The name of the Google Cloud project.
14    project: String,
15    /// The BigQuery dataset.
16    dataset: String,
17    /// The table.
18    table: String,
19}
20
21impl TableName {
22    /// Return the name of the table's project.
23    pub(crate) fn project(&self) -> &str {
24        &self.project
25    }
26
27    /// Return the name of the table's dataset.
28    pub(crate) fn dataset(&self) -> &str {
29        &self.dataset
30    }
31
32    /// Return the bare table name itself, without project or dataset.
33    pub(crate) fn table(&self) -> &str {
34        &self.table
35    }
36
37    /// Return a value which will be formatted as
38    /// `"\`project\`.\`dataset\`.\`table\`"`, with "backtick" quoting.
39    ///
40    /// This form of the name is used in BigQuery "standard SQL".
41    pub(crate) fn dotted_and_quoted(&self) -> DottedTableName {
42        DottedTableName(self)
43    }
44
45    /// Create a temporary table name based on this table name.
46    pub(crate) fn temporary_table_name(
47        &self,
48        temporary_storage: &TemporaryStorage,
49    ) -> Result<TableName> {
50        lazy_static! {
51            static ref DATASET_RE: Regex =
52                Regex::new("^([^:.]+):([^:.]+)$").expect("invalid regex in source");
53        }
54
55        // Decide on what project and dataset to use.
56        let temp = temporary_storage.find_scheme(BigQueryLocator::scheme());
57        let (project, dataset) = if let Some(temp) = temp {
58            // We have a `--temporary=bigquery:...` argument, so extract a project
59            // and dataset name.
60            let cap = DATASET_RE
61                .captures(&temp[BigQueryLocator::scheme().len()..])
62                .ok_or_else(|| {
63                    format_err!("could not parse BigQuery dataset name: {:?}", temp)
64                })?;
65            (cap[1].to_owned(), cap[2].to_owned())
66        } else {
67            // We don't have a `--temporary=bigquery:...` argument, so just pick
68            // something.
69            (self.project.clone(), self.dataset.clone())
70        };
71
72        let tag = TemporaryStorage::random_tag();
73        let table = format!("temp_{}_{}", self.table, tag);
74        Ok(TableName {
75            project,
76            dataset,
77            table,
78        })
79    }
80}
81
82#[test]
83fn temporary_table_name() {
84    let table_name = "project:dataset.table".parse::<TableName>().unwrap();
85
86    // Construct a temporary table name without a `--temporary` argument.
87    let default_temp_name = table_name
88        .temporary_table_name(&TemporaryStorage::new(vec![]))
89        .unwrap()
90        .to_string();
91    assert!(default_temp_name.starts_with("project:dataset.temp_table_"));
92
93    // Now try it with a `--temporary` argument.
94    let temporary_storage =
95        TemporaryStorage::new(vec!["bigquery:project2:temp".to_owned()]);
96    let temp_name = table_name
97        .temporary_table_name(&temporary_storage)
98        .unwrap()
99        .to_string();
100    assert!(temp_name.starts_with("project2:temp.temp_table_"));
101}
102
103impl fmt::Display for TableName {
104    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
105        write!(f, "{}:{}.{}", self.project, self.dataset, self.table)
106    }
107}
108
109impl FromStr for TableName {
110    type Err = Error;
111
112    fn from_str(s: &str) -> Result<Self> {
113        lazy_static! {
114            static ref RE: Regex = Regex::new("^([^:.`]+):([^:.`]+).([^:.`]+)$")
115                .expect("could not parse built-in regex");
116        }
117        let cap = RE.captures(s).ok_or_else(|| {
118            format_err!("could not parse BigQuery table name: {:?}", s)
119        })?;
120        let (project, dataset, table) = (&cap[1], &cap[2], &cap[3]);
121        Ok(TableName {
122            project: project.to_string(),
123            dataset: dataset.to_string(),
124            table: table.to_string(),
125        })
126    }
127}
128
129/// A short-lived wrapped type which displays a BigQuery table name as
130/// `"\`project\`.\`dataset\`.\`table\`"`, with "backtick" quoting.
131///
132/// This form of the name is used in BigQuery "standard SQL".
133pub(crate) struct DottedTableName<'a>(&'a TableName);
134
135impl<'a> fmt::Display for DottedTableName<'a> {
136    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
137        write!(
138            f,
139            "{}.{}.{}",
140            Ident(&self.0.project),
141            Ident(&self.0.dataset),
142            Ident(&self.0.table),
143        )
144    }
145}
146
147/// A BigQuery identifier, for formatting purposes.
148pub(crate) struct Ident<'a>(pub(crate) &'a str);
149
150impl<'a> fmt::Display for Ident<'a> {
151    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
152        if self.0.contains('`') {
153            // We can't output identifiers containing backticks.
154            Err(fmt::Error)
155        } else {
156            write!(f, "`{}`", self.0)
157        }
158    }
159}