1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
use proc_macro::TokenStream;
use std::env;
use std::num::ParseIntError;
use std::path::PathBuf;
use std::str::FromStr;
use proc_macro2::Span;
use quote::quote;
use syn::{LitStr};
use syn::__private::TokenStream2;

use db_up_sql_changelog::ChangelogFile;

/// Represents migration files loaded from a directory
#[derive(Debug, Clone)]
struct MigrationInfo {
    version: u32,
    filename: String,
    name: String,
}

/// Attribute macro for automatically generating a `db_up::MigrationStore`
///
/// The macro takes one required literal string parameter representing the directory containing
/// the migration files. Each file must be named like `V<version>_<name>.sql`, where `<version>`
/// is a valid integer and `<name>` is some name describing what the migration does.
///
/// Example:
/// ```ignore
/// use db_up_codegen::migrations;
///
/// #[migrations("examples/migrations/")]
/// struct Migrations {}
///
/// pub fn main() {
///     let migration_store = Migrations {};
///     println!("migrations: {:?}", migration_store.changelogs());
/// }
/// ```
#[proc_macro_attribute]
pub fn migrations(args: TokenStream, input: TokenStream) -> TokenStream {
    // println!("metadata: {:?}", &args);
    // println!("input:    {:?}", &input);

    let input_clone = input.clone();
    let input_struct = syn::parse_macro_input!(input_clone as syn::ItemStruct);
    // println!("input struct: {:?}", &input_struct);

    let path = if args.is_empty() {
        map_to_crate_root(None)
    } else {
        let migrations_path = syn::parse_macro_input!(args as LitStr).value();
        map_to_crate_root(Some(migrations_path.as_str()))
    };
    println!("migrations path: {:?}", path);

    let migrations = get_migrations(&path)
        .expect("Error while gathering migration file information.");
    println!("migrations: {:?}", &migrations);

    let migration_tokens: Vec<TokenStream2> = migrations.iter()
        .map(|migration| {
            let name = migration.name.as_str();
            let version = migration.version;
            let filename = migration.filename.as_str();
            let file_path = path.clone().join(filename).display().to_string();
            let content = std::fs::read_to_string(file_path.as_str())
                .expect(format!("Could not read migration file: {}", file_path).as_str());

            // just check if the changelog can be loaded correctly:
            let _changelog = ChangelogFile::from_string(version.to_string().as_str(), content.as_str())
                .expect(format!("Migration file is not a valid SQL changelog file: {}", file_path).as_str());

            quote! {
                (#version, #name.to_string(), #content)
            }
        })
        .collect();

    let struct_name = syn::Ident::new(input_struct.ident.to_string().as_str(), Span::call_site());
    // println!("struct_name: {}", &struct_name);
    let result = quote! {
        impl db_up::MigrationStore for #struct_name {
            fn changelogs(&self) -> Vec<db_up::ChangelogFile> {
                use db_up::ChangelogFile;

                let mut result: Vec<ChangelogFile> = [#(#migration_tokens),*].iter()
                .map(|migration| {
                    ChangelogFile::from_string(migration.0.to_string().as_str(), migration.2).unwrap()
                })
                .collect();
                return result;
            }
        }
    };
    // println!("result: {}", result.to_string());

    let input: TokenStream2 = input.into();
    return quote! {
        #input
        #result
    }.into();
}

/// Map a path to the root of the crate
fn map_to_crate_root(path: Option<&str>) -> PathBuf {
    let root = env::var("CARGO_MANIFEST_DIR")
        .map(|root| PathBuf::from(root))
        .expect("Missing CARGO_MANIFEST_DIR environment variable. Cannot obtain crate root.");
    let result = path.map(|path| root.join(PathBuf::from_str(path)
        .expect("Could not parse filename.")))
        .or(Some(root))
        .unwrap();
    return result;
}

/// List migrations contained inside a directory
fn get_migrations(path: &PathBuf) -> Result<Vec<MigrationInfo>, std::io::Error> {
    let result: Vec<MigrationInfo> = std::fs::read_dir(path)?
        .filter(|entry| entry.is_ok())
        .map(|entry| entry.unwrap().file_name().to_str().map(|v| v.to_string()))
        .filter(|filename| filename.is_some())
        .map(|filename| filename.unwrap())
        .filter(|filename| filename.starts_with("V") && filename.ends_with(".sql"))
        .map(|filename| {
            let index = filename.find("_");
            let mut version = "";
            let mut name = "";
            if let Some(index) = index {
                if index > 1 && index < filename.len() - "V.sql".len() {
                    if filename[1..index].chars().all(|ch| ch >= '0' && ch <= '9') {
                        version = &filename[1..index];
                        name = &filename[(index + 1)..(filename.len() - ".sql".len())];
                    }
                }
            }

            return if version.is_empty() {
                None
            } else {
                let result: Result<Option<u32>, ParseIntError> = version.parse::<u32>()
                    .map(|version| Some(version))
                    .or(Ok(None));

                let result = result.unwrap()
                    .map(|version| {
                        MigrationInfo {
                            version,
                            filename: filename.to_string(),
                            name: name.to_string()
                        }
                    });
                return result
            };
        })
        .filter(|info| info.is_some())
        .map(|info| info.unwrap())
        .collect();
    return Ok(result);
}

#[cfg(test)]
mod test {
    #[test]
    pub fn test_get_migrations() {
        let path = crate::map_to_crate_root(Some("examples/migrations"));
        let result = crate::get_migrations(&path);
        match result {
            Ok(migrations) => {
                assert_eq!(migrations.len(), 2, "Two migrations have been successfully loaded.");
            }
            Err(err) => {
                assert!(false, "Migration loading failed: {}", err);
            }
        }
    }
}