<commit_before>use std::error::Error;

use diesel::*;
use diesel::backend::Backend;
use diesel::expression::NonAggregate;
use diesel::query_builder::{QueryId, QueryFragment};
use diesel::types::FromSql;
#[cfg(feature="postgres")]
use diesel::pg::Pg;
#[cfg(feature="mysql")]
use diesel::mysql::Mysql;

use table_data::TableData;
use super::data_structures::*;

pub trait UsesInformationSchema: Backend {
    type TypeColumn: SelectableExpression<
        self::information_schema::columns::table,
        types::Text,
    > + NonAggregate + QueryId + QueryFragment<Self>;

    fn type_column() -> Self::TypeColumn;
}

#[cfg(feature="postgres")]
impl UsesInformationSchema for Pg {
    type TypeColumn = self::information_schema::columns::udt_name;

    fn type_column() -> Self::TypeColumn {
        self::information_schema::columns::udt_name
    }
}

#[cfg(feature="mysql")]
impl UsesInformationSchema for Mysql {
    type TypeColumn = self::information_schema::columns::column_type;

    fn type_column() -> Self::TypeColumn {
        self::information_schema::columns::column_type
    }
}

mod information_schema {
    table! {
        information_schema.tables (table_schema, table_name) {
            table_schema -> VarChar,
            table_name -> VarChar,
            table_type -> VarChar,
        }
    }

    table! {
        information_schema.columns (table_schema, table_name, column_name) {
            table_schema -> VarChar,
            table_name -> VarChar,
            column_name -> VarChar,
            is_nullable -> VarChar,
            ordinal_position -> BigInt,
            udt_name -> VarChar,
            column_type -> VarChar,
        }
    }

    table! {
        information_schema.key_column_usage (table_schema, table_name, column_name, constraint_name) {
            table_schema -> VarChar,
            table_name -> VarChar,
            column_name -> VarChar,
            constraint_name -> VarChar,
            ordinal_position -> BigInt,
        }
    }

    table! {
        information_schema.table_constraints (table_schema, table_name, constraint_name) {
            table_schema -> VarChar,
            table_name -> VarChar,
            constraint_name -> VarChar,
            constraint_type -> VarChar,
        }
    }
}

pub fn determine_column_type(attr: &ColumnInformation) -> Result<ColumnType, Box<Error>> {
    let is_array = attr.type_name.starts_with('_');
    let tpe = if is_array {
        &attr.type_name[1..]
    } else {
        &attr.type_name
    };

    let tpe = if let Some(idx) = tpe.find('(') {
        &tpe[..idx]
    } else {
        tpe
    };

    Ok(ColumnType {
        path: vec!["diesel".into(), "types".into(), capitalize(tpe)],
        is_array: is_array,
        is_nullable: attr.nullable,
    })
}

fn capitalize(name: &str) -> String {
    name[..1].to_uppercase() + &name[1..]
}

pub fn get_table_data<Conn>(conn: &Conn, table: &TableData)
    -> QueryResult<Vec<ColumnInformation>> where
        Conn: Connection,
        Conn::Backend: UsesInformationSchema,
        String: FromSql<types::Text, Conn::Backend>,
{
    use self::information_schema::columns::dsl::*;

    let type_column = <Conn::Backend as UsesInformationSchema>::type_column();
    columns.select((column_name, type_column, is_nullable))
        .filter(table_name.eq(&table.name))
        .filter(table_schema.nullable().eq(&table.schema))
        .order(ordinal_position)
        .load(conn)
}

pub fn get_primary_keys<Conn>(conn: &Conn, table: &TableData)
    -> QueryResult<Vec<String>> where
        Conn: Connection,
        Conn::Backend: UsesInformationSchema,
        String: FromSql<types::Text, Conn::Backend>,
{
    use self::information_schema::table_constraints::{self, constraint_type};
    use self::information_schema::key_column_usage::dsl::*;

    let pk_query = table_constraints::table.select(table_constraints::constraint_name)
        .filter(constraint_type.eq("PRIMARY KEY"));

    key_column_usage.select(column_name)
        .filter(constraint_name.eq_any(pk_query))
        .filter(table_name.eq(&table.name))
        .filter(table_schema.nullable().eq(&table.schema))
        .order(ordinal_position)
        .load(conn)
}

pub fn load_table_names<Conn>(connection: &Conn, schema_name: Option<&str>)
    -> Result<Vec<TableData>, Box<Error>> where
        Conn: Connection,
        Conn::Backend: UsesInformationSchema,
        String: FromSql<types::Text, Conn::Backend>,
{
    use self::information_schema::tables::dsl::*;

    let schema_name = schema_name.unwrap_or("public");

    tables.select((table_name, table_schema))
        .filter(table_schema.eq(schema_name))
        .filter(table_name.not_like("\\_\\_%"))
        .filter(table_type.like("BASE TABLE"))
        .load(connection)
        .map_err(Into::into)
}

#[cfg(all(test, feature="postgres"))]
mod tests {
    extern crate dotenv;

    use diesel::pg::PgConnection;
    use self::dotenv::dotenv;
    use std::env;
    use super::*;

    fn connection() -> PgConnection {
        let _ = dotenv();
<commit_msg>Ensure schema inference loads tables in a deterministic order<commit_after>use std::error::Error;

use diesel::*;
use diesel::backend::Backend;
use diesel::expression::NonAggregate;
use diesel::query_builder::{QueryId, QueryFragment};
use diesel::types::FromSql;
#[cfg(feature="postgres")]
use diesel::pg::Pg;
#[cfg(feature="mysql")]
use diesel::mysql::Mysql;

use table_data::TableData;
use super::data_structures::*;

pub trait UsesInformationSchema: Backend {
    type TypeColumn: SelectableExpression<
        self::information_schema::columns::table,
        types::Text,
    > + NonAggregate + QueryId + QueryFragment<Self>;

    fn type_column() -> Self::TypeColumn;
}

#[cfg(feature="postgres")]
impl UsesInformationSchema for Pg {
    type TypeColumn = self::information_schema::columns::udt_name;

    fn type_column() -> Self::TypeColumn {
        self::information_schema::columns::udt_name
    }
}

#[cfg(feature="mysql")]
impl UsesInformationSchema for Mysql {
    type TypeColumn = self::information_schema::columns::column_type;

    fn type_column() -> Self::TypeColumn {
        self::information_schema::columns::column_type
    }
}

mod information_schema {
    table! {
        information_schema.tables (table_schema, table_name) {
            table_schema -> VarChar,
            table_name -> VarChar,
            table_type -> VarChar,
        }
    }

    table! {
        information_schema.columns (table_schema, table_name, column_name) {
            table_schema -> VarChar,
            table_name -> VarChar,
            column_name -> VarChar,
            is_nullable -> VarChar,
            ordinal_position -> BigInt,
            udt_name -> VarChar,
            column_type -> VarChar,
        }
    }

    table! {
        information_schema.key_column_usage (table_schema, table_name, column_name, constraint_name) {
            table_schema -> VarChar,
            table_name -> VarChar,
            column_name -> VarChar,
            constraint_name -> VarChar,
            ordinal_position -> BigInt,
        }
    }

    table! {
        information_schema.table_constraints (table_schema, table_name, constraint_name) {
            table_schema -> VarChar,
            table_name -> VarChar,
            constraint_name -> VarChar,
            constraint_type -> VarChar,
        }
    }
}

pub fn determine_column_type(attr: &ColumnInformation) -> Result<ColumnType, Box<Error>> {
    let is_array = attr.type_name.starts_with('_');
    let tpe = if is_array {
        &attr.type_name[1..]
    } else {
        &attr.type_name
    };

    let tpe = if let Some(idx) = tpe.find('(') {
        &tpe[..idx]
    } else {
        tpe
    };

    Ok(ColumnType {
        path: vec!["diesel".into(), "types".into(), capitalize(tpe)],
        is_array: is_array,
        is_nullable: attr.nullable,
    })
}

fn capitalize(name: &str) -> String {
    name[..1].to_uppercase() + &name[1..]
}

pub fn get_table_data<Conn>(conn: &Conn, table: &TableData)
    -> QueryResult<Vec<ColumnInformation>> where
        Conn: Connection,
        Conn::Backend: UsesInformationSchema,
        String: FromSql<types::Text, Conn::Backend>,
{
    use self::information_schema::columns::dsl::*;

    let type_column = <Conn::Backend as UsesInformationSchema>::type_column();
    columns.select((column_name, type_column, is_nullable))
        .filter(table_name.eq(&table.name))
        .filter(table_schema.nullable().eq(&table.schema))
        .order(ordinal_position)
        .load(conn)
}

pub fn get_primary_keys<Conn>(conn: &Conn, table: &TableData)
    -> QueryResult<Vec<String>> where
        Conn: Connection,
        Conn::Backend: UsesInformationSchema,
        String: FromSql<types::Text, Conn::Backend>,
{
    use self::information_schema::table_constraints::{self, constraint_type};
    use self::information_schema::key_column_usage::dsl::*;

    let pk_query = table_constraints::table.select(table_constraints::constraint_name)
        .filter(constraint_type.eq("PRIMARY KEY"));

    key_column_usage.select(column_name)
        .filter(constraint_name.eq_any(pk_query))
        .filter(table_name.eq(&table.name))
        .filter(table_schema.nullable().eq(&table.schema))
        .order(ordinal_position)
        .load(conn)
}

pub fn load_table_names<Conn>(connection: &Conn, schema_name: Option<&str>)
    -> Result<Vec<TableData>, Box<Error>> where
        Conn: Connection,
        Conn::Backend: UsesInformationSchema,
        String: FromSql<types::Text, Conn::Backend>,
{
    use self::information_schema::tables::dsl::*;

    let schema_name = schema_name.unwrap_or("public");

    tables.select((table_name, table_schema))
        .filter(table_schema.eq(schema_name))
        .filter(table_name.not_like("\\_\\_%"))
        .filter(table_type.like("BASE TABLE"))
        .order(table_name)
        .load(connection)
        .map_err(Into::into)
}

#[cfg(all(test, feature="postgres"))]
mod tests {
    extern crate dotenv;

    use diesel::pg::PgConnection;
    use self::dotenv::dotenv;
    use std::env;
    use super::*;

    fn connection() -> PgConnection {
        let _ = dotenv();

{'fraction_non_alphanumeric': 0.13490741647888704, 'fraction_numerical': 0.0005878318800822965, 'mean_word_length': 3.0717989629038693, 'pattern_counts': {'":': 0, '<': 41, '<?xml version=': 0, '>': 97, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3097671', 'n_tokens_mistral': 3167, 'n_tokens_neox': 3137, 'n_words': 551}