liboxen 0.53.0

use crate::core::db::data_frames::df_db::with_hardened_query_conn;
use crate::core::df::{sql, tabular};
use crate::core::staged::get_staged_db_manager;
use crate::error::OxenError;
use crate::model::ParsedResource;
use crate::model::data_frame::{DataFrameSchemaSize, DataFrameSlice, DataFrameSliceSchemas};
use crate::model::merkle_tree::node::EMerkleTreeNode;
use crate::model::metadata::generic_metadata::GenericMetadata;
use crate::model::metadata::metadata_tabular::MetadataTabularImpl;
use crate::model::{Commit, DataFrameSize, LocalRepository, Schema};
use crate::opts::DFOpts;
use crate::repositories;
use polars::prelude::IntoLazy as _;

use std::path::Path;

pub mod schemas;

pub async fn get_slice(
    repo: &LocalRepository,
    resource: &ParsedResource,
    path: impl AsRef<Path>,
    opts: &DFOpts,
) -> Result<DataFrameSlice, OxenError> {
    let workspace = resource.workspace.as_ref();
    let commit = match workspace {
        Some(ws) => ws.commit.clone(),
        None => resource
            .commit
            .clone()
            .ok_or_else(|| OxenError::basic_str("Commit not found"))?,
    };

    let (staged_repo, base_repo) = match workspace {
        Some(ws) => (&ws.workspace_repo, repo),
        None => (repo, repo),
    };

    let file_node = match workspace {
        Some(ws) => {
            let staged_db_manager = get_staged_db_manager(staged_repo)?;
            // Try staged DB first
            if let Some(staged_node) = staged_db_manager.read_from_staged_db(&path)? {
                match staged_node.node.node {
                    EMerkleTreeNode::File(f) => Ok(f),
                    _ => Err(OxenError::NotAFile(path.as_ref().to_path_buf().into())),
                }?
            } else {
                // Fall back to commit tree using workspace's commit
                let commit = &ws.commit;
                repositories::tree::get_file_by_path(base_repo, commit, &path)?
                    .ok_or_else(|| OxenError::path_does_not_exist(path.as_ref()))?
            }
        }
        None => {
            let commit = resource
                .commit
                .as_ref()
                .ok_or_else(|| OxenError::basic_str("Commit not found"))?;
            repositories::tree::get_file_by_path(base_repo, commit, &path)?
                .ok_or_else(|| OxenError::path_does_not_exist(path.as_ref()))?
        }
    };

    log::debug!("get_slice file_node {file_node:?}");

    let metadata: Result<MetadataTabularImpl, OxenError> = match file_node.metadata() {
        Some(metadata) => match metadata {
            GenericMetadata::MetadataTabular(metadata) => Ok(metadata.tabular.clone()),
            _ => {
                return Err(OxenError::basic_str("Metadata is not tabular"));
            }
        },
        None => {
            return Err(OxenError::basic_str("File node does not have metadata"));
        }
    };
    let metadata = metadata?;
    log::debug!("get_slice metadata {metadata:?}");

    let source_schema = metadata.schema;
    let data_frame_size = DataFrameSize {
        width: metadata.width,
        height: metadata.height,
    };

    let handle_sql_result = handle_sql_querying(repo, &commit, path, opts, &data_frame_size).await;
    if let Ok(response) = handle_sql_result {
        return Ok(response);
    }
    // Read the data frame from the version path
    let version_store = repo.version_store();
    let df = tabular::read_version_df(
        &version_store,
        &file_node.hash().to_string(),
        file_node.extension(),
        opts,
    )
    .await?;
    log::debug!("get_slice df {:?}", df.height());

    // Check what the view height is
    let view_height = if opts.has_filter_transform() {
        df.height()
    } else {
        data_frame_size.height
    };

    // Update the schema metadata from the source schema
    let mut slice_schema = Schema::from_polars(df.schema());
    slice_schema.update_metadata_from_schema(&source_schema);
    log::debug!("get_slice slice_schema {slice_schema:?}");
    // Return a DataFrameSlice
    Ok(DataFrameSlice {
        schemas: DataFrameSliceSchemas {
            source: DataFrameSchemaSize {
                size: data_frame_size,
                schema: source_schema,
            },
            slice: DataFrameSchemaSize {
                size: DataFrameSize {
                    width: df.width(),
                    height: view_height,
                },
                schema: slice_schema,
            },
        },
        slice: df,
        total_entries: view_height,
    })
}

async fn handle_sql_querying(
    repo: &LocalRepository,
    commit: &Commit,
    path: impl AsRef<Path>,
    opts: &DFOpts,
    data_frame_size: &DataFrameSize,
) -> Result<DataFrameSlice, OxenError> {
    let path = path.as_ref();

    if let Some(sql) = opts.sql.clone() {
        // Finding the queryable workspace opens DuckDB to check the index, and the query reads it.
        // Both are sync DB/filesystem work, so run them as one blocking unit off the worker.
        let query_repo = repo.clone();
        let query_commit = commit.clone();
        let query_path = path.to_path_buf();
        let (workspace, df) = tokio::task::spawn_blocking(move || -> Result<_, OxenError> {
            let workspace =
                crate::core::v_latest::workspaces::data_frames::get_queryable_data_frame_workspace(
                    &query_repo,
                    &query_path,
                    &query_commit,
                )?;
            let db_path =
                repositories::workspaces::data_frames::duckdb_path(&workspace, &query_path);
            let df = with_hardened_query_conn(&db_path, |conn| sql::query_df(conn, sql, None))?;
            Ok((workspace, df))
        })
        .await??;
        log::debug!("handle_sql_querying got df {df:?}");
        let paginated_df = tabular::collect_with_opts(df.clone().lazy(), opts.clone()).await?;

        // Reading the stored schema is a sync repo/tree read; keep it off the worker too.
        let schema_repo = repo.clone();
        let schema_commit = workspace.commit.clone();
        let schema_path = path.to_path_buf();
        let stored_schema = tokio::task::spawn_blocking(move || {
            repositories::data_frames::schemas::get_by_path(
                &schema_repo,
                &schema_commit,
                &schema_path,
            )
        })
        .await??;
        let source_schema =
            stored_schema.unwrap_or_else(|| Schema::from_polars(paginated_df.schema()));

        let mut slice_schema = Schema::from_polars(df.schema());
        slice_schema.update_metadata_from_schema(&source_schema);

        return Ok(DataFrameSlice {
            schemas: DataFrameSliceSchemas {
                source: DataFrameSchemaSize {
                    size: data_frame_size.clone(),
                    schema: source_schema,
                },
                slice: DataFrameSchemaSize {
                    size: DataFrameSize {
                        width: paginated_df.width(),
                        height: paginated_df.height(),
                    },
                    schema: slice_schema,
                },
            },
            slice: paginated_df,
            total_entries: df.height(),
        });
    }

    Err(OxenError::basic_str("Could not query data frame"))
}