nordisk/skade

skade-flight/src/keys.rs

//! Key metadata: `GetPrimaryKeys`, `GetExportedKeys`, `GetImportedKeys`,
//! `GetCrossReference`.
//!
//! arrow-flight 59 ships no builders for these four, so the schemas are
//! written here from FlightSql.proto, verbatim.
//!
//! Iceberg has no foreign keys, so the three FK commands answer an EMPTY
//! result with the right schema (a driver that asks must get the schema it
//! expects, not `UNIMPLEMENTED` — Power BI's navigator asks). Iceberg DOES
//! have **identifier fields** (`schema.identifier-field-ids`, what skade's
//! `create_table_with_identity` sets), and those are the table's logical
//! primary key — so `GetPrimaryKeys` reports them, in declaration order.

use std::sync::{Arc, LazyLock};

use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray, UInt8Array};
use arrow_schema::{DataType, Field, Schema, SchemaRef};

/// `CommandGetPrimaryKeys` result schema (FlightSql.proto).
pub static PRIMARY_KEYS_SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
    Arc::new(Schema::new(vec![
        Field::new("catalog_name", DataType::Utf8, true),
        Field::new("db_schema_name", DataType::Utf8, true),
        Field::new("table_name", DataType::Utf8, false),
        Field::new("column_name", DataType::Utf8, false),
        Field::new("key_name", DataType::Utf8, true),
        Field::new("key_sequence", DataType::Int32, false),
    ]))
});

/// `CommandGetExportedKeys` / `ImportedKeys` / `CrossReference` result schema.
pub static FOREIGN_KEYS_SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
    Arc::new(Schema::new(vec![
        Field::new("pk_catalog_name", DataType::Utf8, true),
        Field::new("pk_db_schema_name", DataType::Utf8, true),
        Field::new("pk_table_name", DataType::Utf8, false),
        Field::new("pk_column_name", DataType::Utf8, false),
        Field::new("fk_catalog_name", DataType::Utf8, true),
        Field::new("fk_db_schema_name", DataType::Utf8, true),
        Field::new("fk_table_name", DataType::Utf8, false),
        Field::new("fk_column_name", DataType::Utf8, false),
        Field::new("key_sequence", DataType::Int32, false),
        Field::new("fk_key_name", DataType::Utf8, true),
        Field::new("pk_key_name", DataType::Utf8, true),
        Field::new("update_rule", DataType::UInt8, false),
        Field::new("delete_rule", DataType::UInt8, false),
    ]))
});

/// One primary-key column: `(column_name, key_sequence)`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PkColumn {
    pub column: String,
    pub sequence: i32,
}

/// Build the `GetPrimaryKeys` batch for one table.
pub fn primary_keys(catalog: &str, schema: &str, table: &str, columns: &[PkColumn]) -> RecordBatch {
    let n = columns.len();
    let cols: Vec<ArrayRef> = vec![
        Arc::new(StringArray::from(vec![Some(catalog); n])),
        Arc::new(StringArray::from(vec![Some(schema); n])),
        Arc::new(StringArray::from(vec![table; n])),
        Arc::new(StringArray::from(columns.iter().map(|c| c.column.as_str()).collect::<Vec<_>>())),
        Arc::new(StringArray::from(vec![Some("identifier"); n])),
        Arc::new(Int32Array::from(columns.iter().map(|c| c.sequence).collect::<Vec<_>>())),
    ];
    RecordBatch::try_new(PRIMARY_KEYS_SCHEMA.clone(), cols).expect("pk batch matches its schema")
}

/// The empty foreign-key batch (Iceberg has no FK constraints).
pub fn no_foreign_keys() -> RecordBatch {
    let s = || Arc::new(StringArray::new_null(0)) as ArrayRef;
    RecordBatch::try_new(
        FOREIGN_KEYS_SCHEMA.clone(),
        vec![
            s(),
            s(),
            s(),
            s(),
            s(),
            s(),
            s(),
            s(),
            Arc::new(Int32Array::from(Vec::<i32>::new())),
            s(),
            s(),
            Arc::new(UInt8Array::from(Vec::<u8>::new())),
            Arc::new(UInt8Array::from(Vec::<u8>::new())),
        ],
    )
    .expect("empty fk batch matches its schema")
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn pk_batch_shape() {
        let b = primary_keys(
            "skade",
            "main",
            "people",
            &[PkColumn { column: "id".into(), sequence: 1 }, PkColumn { column: "org".into(), sequence: 2 }],
        );
        assert_eq!(b.num_rows(), 2);
        assert_eq!(b.schema(), *PRIMARY_KEYS_SCHEMA);
        let empty = primary_keys("skade", "main", "people", &[]);
        assert_eq!(empty.num_rows(), 0);
    }

    #[test]
    fn fk_batch_is_empty_with_spec_schema() {
        let b = no_foreign_keys();
        assert_eq!(b.num_rows(), 0);
        assert_eq!(b.schema().fields().len(), 13);
        assert_eq!(b.schema().field(11).data_type(), &DataType::UInt8);
    }
}