Skip to main content

zql_cli/db/plugin/
tsql.rs

1use crate::core::catalog::Catalog;
2use crate::db::context::Context;
3use crate::db::database::{Database, Plugin};
4use crate::error::MyResult;
5use crate::util::convert::{parse_null, parse_type};
6use itertools::Itertools;
7use std::borrow::Cow;
8
9pub struct TSQLPlugin;
10
11impl TSQLPlugin {
12    pub fn new() -> Self {
13        Self { }
14    }
15}
16
17impl Plugin for TSQLPlugin {
18    fn populate_catalog(&self, database: &Database, catalog: &mut Catalog) -> MyResult<()> {
19        fetch_databases(database, catalog)?;
20        fetch_keywords(catalog);
21        catalog.insert_keyword(Cow::Borrowed(self.get_batch()));
22        Ok(())
23    }
24
25    fn query_context(&self, database: &Database) -> Context {
26        let (name, schema) = fetch_context(database);
27        Context::new(Some(name), Some(schema))
28    }
29
30    fn get_batch(&self) -> &str {
31        "GO"
32    }
33}
34
35// [https://learn.microsoft.com/en-us/sql/t-sql/language-elements/reserved-keywords-transact-sql]
36const TSQL_KEYWORDS: &[&str] = &[
37    "ADD", "ALL", "ALTER", "AND", "ANY", "AS", "ASC", "AUTHORIZATION",
38    "BACKUP", "BEGIN", "BETWEEN", "BREAK", "BROWSE", "BULK", "BY",
39    "CASCADE", "CASE", "CHECK", "CHECKPOINT", "CLOSE", "CLUSTERED",
40    "COALESCE", "COLLATE", "COLUMN", "COMMIT", "COMPUTE", "CONSTRAINT",
41    "CONTAINS", "CONTAINSTABLE", "CONTINUE", "CONVERT", "CREATE", "CROSS",
42    "CURRENT", "CURRENT_DATE", "CURRENT_TIME", "CURRENT_TIMESTAMP",
43    "CURRENT_USER", "CURSOR", "DATABASE", "DBCC", "DEALLOCATE", "DECLARE",
44    "DEFAULT", "DELETE", "DENY", "DESC", "DISK", "DISTINCT", "DISTRIBUTED",
45    "DOUBLE", "DROP", "DUMP", "ELSE", "END", "ERRLVL", "ESCAPE", "EXCEPT",
46    "EXEC", "EXECUTE", "EXISTS", "EXIT", "EXTERNAL", "FETCH", "FILE",
47    "FILLFACTOR", "FOR", "FOREIGN", "FREETEXT", "FREETEXTTABLE", "FROM",
48    "FULL", "FUNCTION", "GOTO", "GRANT", "GROUP", "HAVING", "HOLDLOCK",
49    "IDENTITY", "IDENTITYCOL", "IDENTITY_INSERT", "IF", "IN", "INDEX",
50    "INNER", "INSERT", "INTERSECT", "INTO", "IS", "JOIN", "KEY", "KILL",
51    "LEFT", "LIKE", "LINENO", "LOAD", "MERGE", "NATIONAL", "NOCHECK",
52    "NONCLUSTERED", "NOT", "NULL", "NULLIF", "OF", "OFF", "OFFSETS",
53    "ON", "OPEN", "OPENDATASOURCE", "OPENQUERY", "OPENROWSET", "OPENXML",
54    "OPTION", "OR", "ORDER", "OUTER", "OVER", "PERCENT", "PIVOT", "PLAN",
55    "PRECISION", "PRIMARY", "PRINT", "PROC", "PROCEDURE", "PUBLIC",
56    "RAISERROR", "READ", "READTEXT", "RECONFIGURE", "REFERENCES",
57    "REPLICATION", "RESTORE", "RESTRICT", "RETURN", "REVERT", "REVOKE",
58    "RIGHT", "ROLLBACK", "ROWCOUNT", "ROWGUIDCOL", "RULE", "SAVE",
59    "SCHEMA", "SECURITYAUDIT", "SELECT", "SEMANTICKEYPHRASETABLE",
60    "SEMANTICSIMILARITYDETAILSTABLE", "SEMANTICSIMILARITYTABLE",
61    "SESSION_USER", "SET", "SETUSER", "SHUTDOWN", "SOME", "STATISTICS",
62    "SYSTEM_USER", "TABLE", "TABLESAMPLE", "TEXTSIZE", "THEN", "TO",
63    "TOP", "TRAN", "TRANSACTION", "TRIGGER", "TRUNCATE", "TRY_CONVERT",
64    "TSEQUAL", "UNION", "UNIQUE", "UNPIVOT", "UPDATE", "UPDATETEXT",
65    "USE", "USER", "VALUES", "VARYING", "VIEW", "WAITFOR", "WHEN",
66    "WHERE", "WHILE", "WITH", "WITHIN", "GROUP", "WRITETEXT",
67];
68
69fn fetch_databases(database: &Database, catalog: &mut Catalog) -> MyResult<()> {
70    database.fetch_one("SELECT name FROM sys.databases", |name| {
71        fetch_columns(database, catalog, name.clone())?;
72        fetch_procedures(database, catalog, name)?;
73        Ok(())
74    })?;
75    Ok(())
76}
77
78fn fetch_columns(database: &Database, catalog: &mut Catalog, name: Cow<str>) -> MyResult<()> {
79    let query = format!("\
80SELECT table_schema, table_name, column_name, ordinal_position, is_nullable,
81data_type, character_maximum_length, numeric_precision, numeric_scale
82FROM {}.information_schema.columns", name);
83    database.fetch_records(&query, |record| {
84        if let Some((schema, table, column, index, null, dtype, dsize, nsize, nprec)) = record.into_iter().next_tuple() {
85            let dsize = nsize.parse().unwrap_or_else(|_| dsize.parse().unwrap_or_default());
86            let dprec = nprec.parse().unwrap_or_default();
87            let dtype = parse_type(dtype, dsize, dprec);
88            let null = parse_null(null);
89            let index = index.parse().unwrap_or_default();
90            catalog.insert_column(
91                name.clone(),
92                schema,
93                table,
94                column,
95                dtype,
96                null,
97                index,
98            );
99        }
100        Ok(())
101    })?;
102    Ok(())
103}
104
105fn fetch_procedures(database: &Database, catalog: &mut Catalog, name: Cow<str>) -> MyResult<()> {
106    let query = format!("\
107SELECT routine_schema, routine_name
108FROM {}.information_schema.routines", name);
109    database.fetch_two(&query, |schema, routine| {
110        catalog.insert_procedure(name.clone(), schema, routine);
111        Ok(())
112    })?;
113    Ok(())
114}
115
116fn fetch_keywords(catalog: &mut Catalog) {
117    for keyword in TSQL_KEYWORDS {
118        catalog.insert_keyword(Cow::Borrowed(keyword));
119    }
120}
121
122fn fetch_context(database: &Database) -> (String, String) {
123    let mut name = String::new();
124    let mut schema = String::new();
125    let _ = database.fetch_one("SELECT DB_NAME()", |value| {
126        name = value.to_string();
127        Ok(())
128    });
129    let _ = database.fetch_one("SELECT SCHEMA_NAME()", |value| {
130        schema = value.to_string();
131        Ok(())
132    });
133    (name, schema)
134}