linter: add fix for prefer-timestamptz (#617)

Author: sbdchd

crates/squawk_linter/src/lib.rs

@@ -246,6 +246,12 @@ impl Edit {
             text: Some(text.into()),
         }
     }
+    pub fn replace<T: Into<String>>(text_range: TextRange, text: T) -> Self {
+        Self {
+            text_range,
+            text: Some(text.into()),
+        }
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]

crates/squawk_linter/src/rules/prefer_robust_stmts.rs

@@ -242,48 +242,12 @@ mod test {
     use insta::{assert_debug_snapshot, assert_snapshot};
 
     use crate::{
-        Edit, Linter, Rule,
-        test_utils::{lint, lint_with_assume_in_transaction},
+        Rule,
+        test_utils::{fix_sql, lint, lint_with_assume_in_transaction},
     };
 
     fn fix(sql: &str) -> String {
-        let file = squawk_syntax::SourceFile::parse(sql);
-        assert_eq!(file.errors().len(), 0);
-        assert_eq!(file.errors().len(), 0, "Shouldn't start with syntax errors");
-        let mut linter = Linter::from([Rule::PreferRobustStmts]);
-        let errors = linter.lint(&file, sql);
-        assert!(!errors.is_empty(), "Should start with linter errors");
-
-        let fixes = errors.into_iter().flat_map(|x| x.fix).collect::<Vec<_>>();
-
-        let mut result = sql.to_string();
-
-        let mut all_edits: Vec<&Edit> = fixes.iter().flat_map(|fix| &fix.edits).collect();
-
-        all_edits.sort_by(|a, b| b.text_range.start().cmp(&a.text_range.start()));
-
-        for edit in all_edits {
-            let start: usize = edit.text_range.start().into();
-            let end: usize = edit.text_range.end().into();
-            let text = edit.text.as_ref().map_or("", |v| v);
-            result.replace_range(start..end, text);
-        }
-
-        let file = squawk_syntax::SourceFile::parse(&result);
-        assert_eq!(
-            file.errors().len(),
-            0,
-            "Shouldn't introduce any syntax errors"
-        );
-        let mut linter = Linter::from([Rule::PreferRobustStmts]);
-        let errors = linter.lint(&file, &result);
-        assert_eq!(
-            errors.len(),
-            0,
-            "Fixes should remove all the linter errors."
-        );
-
-        result
+        fix_sql(sql, Rule::PreferRobustStmts)
     }
 
     #[test]

crates/squawk_linter/src/rules/prefer_timestamptz.rs

@@ -3,7 +3,7 @@ use squawk_syntax::{
     ast::{self, AstNode},
 };
 
-use crate::{Linter, Rule, Violation};
+use crate::{Edit, Fix, Linter, Rule, Violation};
 use crate::{identifier::Identifier, visitors::check_not_allowed_types};
 
 pub fn is_not_allowed_timestamp(ty: &ast::Type) -> bool {
@@ -44,14 +44,33 @@ pub fn is_not_allowed_timestamp(ty: &ast::Type) -> bool {
     }
 }
 
+fn fix_timestamp(ty: &ast::Type) -> Option<Fix> {
+    match ty {
+        ast::Type::TimeType(_) => {
+            let range = ty.syntax().text_range();
+            let edit = Edit::replace(range, "timestamptz");
+            Some(Fix::new("Replace with `timestamptz`", vec![edit]))
+        }
+        ast::Type::ArrayType(array_type) => {
+            if let Some(inner_ty) = array_type.ty() {
+                fix_timestamp(&inner_ty)
+            } else {
+                None
+            }
+        }
+        _ => None,
+    }
+}
+
 fn check_ty_for_timestamp(ctx: &mut Linter, ty: Option<ast::Type>) {
     if let Some(ty) = ty {
         if is_not_allowed_timestamp(&ty) {
+            let fix = fix_timestamp(&ty);
             ctx.report(Violation::for_node(
                 Rule::PreferTimestampTz,
             "When Postgres stores a datetime in a `timestamp` field, Postgres drops the UTC offset. This means 2019-10-11 21:11:24+02 and 2019-10-11 21:11:24-06 will both be stored as 2019-10-11 21:11:24 in the database, even though they are eight hours apart in time.".into(),
                 ty.syntax(),
-            ).help("Use timestamptz instead of timestamp for your column type."));
+            ).help("Use `timestamptz` instead of `timestamp` for your column type.").fix(fix));
         };
     }
 }
@@ -63,10 +82,70 @@ pub(crate) fn prefer_timestamptz(ctx: &mut Linter, parse: &Parse<SourceFile>) {
 
 #[cfg(test)]
 mod test {
-    use insta::assert_debug_snapshot;
+    use insta::{assert_debug_snapshot, assert_snapshot};
 
     use crate::Rule;
-    use crate::test_utils::lint;
+    use crate::test_utils::{fix_sql, lint};
+
+    fn fix(sql: &str) -> String {
+        fix_sql(sql, Rule::PreferTimestampTz)
+    }
+
+    #[test]
+    fn fix_timestamp_to_timestamptz() {
+        assert_snapshot!(fix("
+create table app.users
+(
+    created_ts   timestamp
+);
+"), @r"
+        create table app.users
+        (
+            created_ts   timestamptz
+        );
+        ");
+    }
+
+    #[test]
+    fn fix_timestamp_without_time_zone() {
+        assert_snapshot!(fix("
+create table app.accounts
+(
+    created_ts timestamp without time zone
+);
+"), @r"
+        create table app.accounts
+        (
+            created_ts timestamptz
+        );
+        ");
+    }
+
+    #[test]
+    fn fix_alter_table_timestamp() {
+        assert_snapshot!(fix("
+alter table app.users
+    alter column created_ts type timestamp;
+"), @r"
+        alter table app.users
+            alter column created_ts type timestamptz;
+        ");
+    }
+
+    #[test]
+    fn fix_timestamp_array() {
+        assert_snapshot!(fix("
+create table app.events
+(
+    timestamps timestamp[]
+);
+"), @r"
+        create table app.events
+        (
+            timestamps timestamptz[]
+        );
+        ");
+    }
 
     #[test]
     fn create_table_with_timestamp_err() {

crates/squawk_linter/src/rules/snapshots/squawk_linter__rules__prefer_timestamptz__test__alter_table_with_timestamp_err.snap

@@ -8,17 +8,41 @@ expression: errors
         message: "When Postgres stores a datetime in a `timestamp` field, Postgres drops the UTC offset. This means 2019-10-11 21:11:24+02 and 2019-10-11 21:11:24-06 will both be stored as 2019-10-11 21:11:24 in the database, even though they are eight hours apart in time.",
         text_range: 56..65,
         help: Some(
-            "Use timestamptz instead of timestamp for your column type.",
+            "Use `timestamptz` instead of `timestamp` for your column type.",
+        ),
+        fix: Some(
+            Fix {
+                title: "Replace with `timestamptz`",
+                edits: [
+                    Edit {
+                        text_range: 56..65,
+                        text: Some(
+                            "timestamptz",
+                        ),
+                    },
+                ],
+            },
         ),
-        fix: None,
     },
     Violation {
         code: PreferTimestampTz,
         message: "When Postgres stores a datetime in a `timestamp` field, Postgres drops the UTC offset. This means 2019-10-11 21:11:24+02 and 2019-10-11 21:11:24-06 will both be stored as 2019-10-11 21:11:24 in the database, even though they are eight hours apart in time.",
         text_range: 125..152,
         help: Some(
-            "Use timestamptz instead of timestamp for your column type.",
+            "Use `timestamptz` instead of `timestamp` for your column type.",
+        ),
+        fix: Some(
+            Fix {
+                title: "Replace with `timestamptz`",
+                edits: [
+                    Edit {
+                        text_range: 125..152,
+                        text: Some(
+                            "timestamptz",
+                        ),
+                    },
+                ],
+            },
         ),
-        fix: None,
     },
 ]

crates/squawk_linter/src/rules/snapshots/squawk_linter__rules__prefer_timestamptz__test__create_table_with_timestamp_err.snap

@@ -8,17 +8,41 @@ expression: errors
         message: "When Postgres stores a datetime in a `timestamp` field, Postgres drops the UTC offset. This means 2019-10-11 21:11:24+02 and 2019-10-11 21:11:24-06 will both be stored as 2019-10-11 21:11:24 in the database, even though they are eight hours apart in time.",
         text_range: 43..52,
         help: Some(
-            "Use timestamptz instead of timestamp for your column type.",
+            "Use `timestamptz` instead of `timestamp` for your column type.",
+        ),
+        fix: Some(
+            Fix {
+                title: "Replace with `timestamptz`",
+                edits: [
+                    Edit {
+                        text_range: 43..52,
+                        text: Some(
+                            "timestamptz",
+                        ),
+                    },
+                ],
+            },
         ),
-        fix: None,
     },
     Violation {
         code: PreferTimestampTz,
         message: "When Postgres stores a datetime in a `timestamp` field, Postgres drops the UTC offset. This means 2019-10-11 21:11:24+02 and 2019-10-11 21:11:24-06 will both be stored as 2019-10-11 21:11:24 in the database, even though they are eight hours apart in time.",
         text_range: 99..126,
         help: Some(
-            "Use timestamptz instead of timestamp for your column type.",
+            "Use `timestamptz` instead of `timestamp` for your column type.",
+        ),
+        fix: Some(
+            Fix {
+                title: "Replace with `timestamptz`",
+                edits: [
+                    Edit {
+                        text_range: 99..126,
+                        text: Some(
+                            "timestamptz",
+                        ),
+                    },
+                ],
+            },
         ),
-        fix: None,
     },
 ]

crates/squawk_linter/src/test_utils.rs

@@ -1,4 +1,4 @@
-use crate::{Linter, Rule, Violation};
+use crate::{Edit, Linter, Rule, Violation};
 
 pub(crate) fn lint(sql: &str, rule: Rule) -> Vec<Violation> {
     let file = squawk_syntax::SourceFile::parse(sql);
@@ -14,3 +14,42 @@ pub(crate) fn lint_with_assume_in_transaction(sql: &str, rule: Rule) -> Vec<Viol
     linter.settings.assume_in_transaction = true;
     linter.lint(&file, sql)
 }
+
+pub(crate) fn fix_sql(sql: &str, rule: Rule) -> String {
+    let file = squawk_syntax::SourceFile::parse(sql);
+    assert_eq!(file.errors().len(), 0, "Shouldn't start with syntax errors");
+    let mut linter = Linter::from([rule]);
+    let errors = linter.lint(&file, sql);
+    assert!(!errors.is_empty(), "Should start with linter errors");
+
+    let fixes = errors.into_iter().flat_map(|x| x.fix).collect::<Vec<_>>();
+
+    let mut result = sql.to_string();
+
+    let mut all_edits: Vec<&Edit> = fixes.iter().flat_map(|fix| &fix.edits).collect();
+
+    all_edits.sort_by(|a, b| b.text_range.start().cmp(&a.text_range.start()));
+
+    for edit in all_edits {
+        let start: usize = edit.text_range.start().into();
+        let end: usize = edit.text_range.end().into();
+        let text = edit.text.as_ref().map_or("", |v| v);
+        result.replace_range(start..end, text);
+    }
+
+    let file = squawk_syntax::SourceFile::parse(&result);
+    assert_eq!(
+        file.errors().len(),
+        0,
+        "Shouldn't introduce any syntax errors"
+    );
+    let mut linter = Linter::from([rule]);
+    let errors = linter.lint(&file, &result);
+    assert_eq!(
+        errors.len(),
+        0,
+        "Fixes should remove all the linter errors."
+    );
+
+    result
+}