@lix-js/sdk 0.6.0-preview.0 → 0.6.0-preview.2

package/dist-engine-src/src/session/switch_version.rs

@@ -0,0 +1,121 @@
+use std::sync::Arc;
+
+use serde_json::json;
+
+use crate::transaction::types::StageRow;
+use crate::version::VersionRefReader;
+use crate::LixError;
+use crate::GLOBAL_VERSION_ID;
+
+use super::context::{SessionContext, SessionMode, WORKSPACE_VERSION_KEY};
+
+const KEY_VALUE_SCHEMA_KEY: &str = "lix_key_value";
+const KEY_VALUE_SCHEMA_VERSION: &str = "1";
+
+/// Options for switching a session to another version.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct SwitchVersionOptions {
+    pub version_id: String,
+}
+
+/// Receipt returned after switching to another version.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct SwitchVersionReceipt {
+    pub version_id: String,
+}
+
+impl SessionContext {
+    /// Switches the session's active version selector.
+    ///
+    /// Pinned sessions switch in memory and return a new pinned session.
+    /// Workspace sessions update the shared workspace selector so other
+    /// workspace sessions observe the new active version on their next use.
+    pub async fn switch_version(
+        &self,
+        options: SwitchVersionOptions,
+    ) -> Result<(SessionContext, SwitchVersionReceipt), LixError> {
+        let version_id = options.version_id;
+        let receipt_version_id = version_id.clone();
+        let current_mode = self.mode.clone();
+        let next_mode = self
+            .with_write_transaction(|transaction| {
+                Box::pin(async move {
+                    let head = {
+                        let reader = transaction.version_ref_reader();
+                        reader.load_head_commit_id(&version_id).await?
+                    };
+                    if head.is_none() {
+                        return Err(LixError::version_not_found(
+                            version_id.clone(),
+                            "switch_version",
+                            "target",
+                        ));
+                    }
+
+                    match current_mode {
+                        SessionMode::Pinned { .. } => Ok(SessionMode::Pinned {
+                            version_id: version_id.clone(),
+                        }),
+                        SessionMode::Workspace => {
+                            transaction
+                                .stage_rows(vec![workspace_version_stage_row(&version_id)?])
+                                .await?;
+                            Ok(SessionMode::Workspace)
+                        }
+                    }
+                })
+            })
+            .await?;
+
+        let session = SessionContext::new_with_closed(
+            next_mode,
+            self.storage.clone(),
+            Arc::clone(&self.live_state),
+            Arc::clone(&self.tracked_state),
+            Arc::clone(&self.binary_cas),
+            Arc::clone(&self.changelog),
+            Arc::clone(&self.version_ctx),
+            Arc::clone(&self.schema_registry),
+            self.closed_flag(),
+        );
+        Ok((
+            session,
+            SwitchVersionReceipt {
+                version_id: receipt_version_id,
+            },
+        ))
+    }
+}
+
+fn workspace_version_stage_row(version_id: &str) -> Result<StageRow, LixError> {
+    Ok(StageRow {
+        entity_id: Some(crate::entity_identity::EntityIdentity::single(
+            WORKSPACE_VERSION_KEY,
+        )),
+        schema_key: KEY_VALUE_SCHEMA_KEY.to_string(),
+        file_id: None,
+        snapshot_content: Some(encode_snapshot(json!({
+            "key": WORKSPACE_VERSION_KEY,
+            "value": version_id,
+        }))?),
+        metadata: None,
+        origin: None,
+        schema_version: KEY_VALUE_SCHEMA_VERSION.to_string(),
+        created_at: None,
+        updated_at: None,
+        global: true,
+        change_id: None,
+        commit_id: None,
+        untracked: true,
+        version_id: GLOBAL_VERSION_ID.to_string(),
+    })
+}
+
+fn encode_snapshot(value: serde_json::Value) -> Result<String, LixError> {
+    serde_json::to_string(&value).map_err(|error| {
+        LixError::new(
+            "LIX_ERROR_UNKNOWN",
+            format!("engine2 switch_version snapshot serialization failed: {error}"),
+        )
+    })
+}

package/dist-engine-src/src/sql2/change_provider.rs

@@ -0,0 +1,337 @@
+use std::any::Any;
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use datafusion::arrow::array::{ArrayRef, StringArray};
+use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use datafusion::arrow::record_batch::RecordBatch;
+use datafusion::catalog::{Session, TableProvider};
+use datafusion::common::{DataFusionError, Result};
+use datafusion::datasource::TableType;
+use datafusion::execution::TaskContext;
+use datafusion::logical_expr::{Expr, TableProviderFilterPushDown};
+use datafusion::physical_expr::EquivalenceProperties;
+use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType, PlanProperties};
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion::physical_plan::{
+    DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream,
+};
+use futures_util::stream;
+
+use crate::changelog::{materialize_change, ChangelogScanRequest, MaterializedCanonicalChange};
+use crate::serialize_row_metadata;
+use crate::LixError;
+
+use super::record_batch::record_batch_with_row_count;
+use super::result_metadata::json_field;
+use super::SqlChangelogQuerySource;
+
+pub(crate) async fn register_lix_change_provider(
+    session: &datafusion::prelude::SessionContext,
+    query_source: SqlChangelogQuerySource,
+) -> Result<(), LixError> {
+    session
+        .register_table("lix_change", Arc::new(LixChangeProvider::new(query_source)))
+        .map_err(datafusion_error_to_lix_error)?;
+    Ok(())
+}
+
+struct LixChangeProvider {
+    schema: SchemaRef,
+    query_source: SqlChangelogQuerySource,
+}
+
+impl std::fmt::Debug for LixChangeProvider {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("LixChangeProvider").finish()
+    }
+}
+
+impl LixChangeProvider {
+    fn new(query_source: SqlChangelogQuerySource) -> Self {
+        Self {
+            schema: lix_change_schema(),
+            query_source,
+        }
+    }
+}
+
+#[async_trait]
+impl TableProvider for LixChangeProvider {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn schema(&self) -> SchemaRef {
+        Arc::clone(&self.schema)
+    }
+
+    fn table_type(&self) -> TableType {
+        TableType::Base
+    }
+
+    fn supports_filters_pushdown(
+        &self,
+        filters: &[&Expr],
+    ) -> Result<Vec<TableProviderFilterPushDown>> {
+        Ok(filters
+            .iter()
+            .map(|_| TableProviderFilterPushDown::Unsupported)
+            .collect())
+    }
+
+    async fn scan(
+        &self,
+        _state: &dyn Session,
+        projection: Option<&Vec<usize>>,
+        _filters: &[Expr],
+        limit: Option<usize>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        Ok(Arc::new(LixChangeScanExec::new(
+            self.query_source.clone(),
+            projected_schema(&self.schema, projection),
+            projection.cloned(),
+            limit,
+        )))
+    }
+}
+
+struct LixChangeScanExec {
+    query_source: SqlChangelogQuerySource,
+    schema: SchemaRef,
+    projection: Option<Vec<usize>>,
+    limit: Option<usize>,
+    properties: Arc<PlanProperties>,
+}
+
+impl std::fmt::Debug for LixChangeScanExec {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("LixChangeScanExec").finish()
+    }
+}
+
+impl LixChangeScanExec {
+    fn new(
+        query_source: SqlChangelogQuerySource,
+        schema: SchemaRef,
+        projection: Option<Vec<usize>>,
+        limit: Option<usize>,
+    ) -> Self {
+        let properties = PlanProperties::new(
+            EquivalenceProperties::new(schema.clone()),
+            Partitioning::UnknownPartitioning(1),
+            EmissionType::Incremental,
+            Boundedness::Bounded,
+        );
+        Self {
+            query_source,
+            schema,
+            projection,
+            limit,
+            properties: Arc::new(properties),
+        }
+    }
+}
+
+impl DisplayAs for LixChangeScanExec {
+    fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match t {
+            DisplayFormatType::Default | DisplayFormatType::Verbose => {
+                write!(f, "LixChangeScanExec")
+            }
+            DisplayFormatType::TreeRender => write!(f, "LixChangeScanExec"),
+        }
+    }
+}
+
+impl ExecutionPlan for LixChangeScanExec {
+    fn name(&self) -> &str {
+        "LixChangeScanExec"
+    }
+
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn properties(&self) -> &Arc<PlanProperties> {
+        &self.properties
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        Vec::new()
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        if !children.is_empty() {
+            return Err(DataFusionError::Execution(
+                "LixChangeScanExec does not accept children".to_string(),
+            ));
+        }
+        Ok(self)
+    }
+
+    fn execute(
+        &self,
+        partition: usize,
+        _context: Arc<TaskContext>,
+    ) -> Result<SendableRecordBatchStream> {
+        if partition != 0 {
+            return Err(DataFusionError::Execution(format!(
+                "LixChangeScanExec only exposes one partition, got {partition}"
+            )));
+        }
+
+        let query_source = self.query_source.clone();
+        let projection = change_projection_for_scan(self.projection.as_ref());
+        let limit = self.limit;
+        let schema = Arc::clone(&self.schema);
+        let stream = stream::once(async move {
+            let mut json_reader = query_source.json_reader;
+            let canonical_changes = query_source
+                .changelog_reader
+                .scan_changes(&ChangelogScanRequest { limit })
+                .await
+                .map_err(lix_error_to_datafusion_error)?;
+            let mut changes = Vec::with_capacity(canonical_changes.len());
+            for change in canonical_changes {
+                changes.push(
+                    materialize_change(&mut json_reader, change)
+                        .await
+                        .map_err(lix_error_to_datafusion_error)?,
+                );
+            }
+            change_record_batch(&projection, &changes)
+        });
+        Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream)))
+    }
+}
+
+#[derive(Debug, Clone, Copy)]
+enum ChangeColumn {
+    Id,
+    EntityId,
+    SchemaKey,
+    SchemaVersion,
+    FileId,
+    Metadata,
+    CreatedAt,
+    SnapshotContent,
+}
+
+fn lix_change_schema() -> SchemaRef {
+    Arc::new(Schema::new(vec![
+        Field::new("id", DataType::Utf8, false),
+        Field::new("entity_id", DataType::Utf8, false),
+        Field::new("schema_key", DataType::Utf8, false),
+        Field::new("schema_version", DataType::Utf8, false),
+        Field::new("file_id", DataType::Utf8, true),
+        json_field("metadata", true),
+        Field::new("created_at", DataType::Utf8, false),
+        json_field("snapshot_content", true),
+    ]))
+}
+
+fn change_projection_for_scan(projection: Option<&Vec<usize>>) -> Vec<ChangeColumn> {
+    let all_columns = vec![
+        ChangeColumn::Id,
+        ChangeColumn::EntityId,
+        ChangeColumn::SchemaKey,
+        ChangeColumn::SchemaVersion,
+        ChangeColumn::FileId,
+        ChangeColumn::Metadata,
+        ChangeColumn::CreatedAt,
+        ChangeColumn::SnapshotContent,
+    ];
+    projection.map_or(all_columns.clone(), |indices| {
+        indices
+            .iter()
+            .filter_map(|index| all_columns.get(*index).copied())
+            .collect()
+    })
+}
+
+fn projected_schema(schema: &SchemaRef, projection: Option<&Vec<usize>>) -> SchemaRef {
+    match projection {
+        Some(projection) => Arc::new(schema.project(projection).expect("projection is valid")),
+        None => Arc::clone(schema),
+    }
+}
+
+fn change_record_batch(
+    projection: &[ChangeColumn],
+    changes: &[MaterializedCanonicalChange],
+) -> Result<RecordBatch> {
+    let arrays = projection
+        .iter()
+        .map(|column| match column {
+            ChangeColumn::Id => string_array(changes.iter().map(|row| Some(row.id.as_str()))),
+            ChangeColumn::EntityId => Arc::new(StringArray::from(
+                changes
+                    .iter()
+                    .map(|row| {
+                        Some(
+                            row.entity_id
+                                .as_string()
+                                .expect("canonical change entity identity should project"),
+                        )
+                    })
+                    .collect::<Vec<_>>(),
+            )) as ArrayRef,
+            ChangeColumn::SchemaKey => {
+                string_array(changes.iter().map(|row| Some(row.schema_key.as_str())))
+            }
+            ChangeColumn::SchemaVersion => {
+                string_array(changes.iter().map(|row| Some(row.schema_version.as_str())))
+            }
+            ChangeColumn::FileId => string_array(changes.iter().map(|row| row.file_id.as_deref())),
+            ChangeColumn::Metadata => Arc::new(StringArray::from(
+                changes
+                    .iter()
+                    .map(|row| row.metadata.as_ref().map(serialize_row_metadata))
+                    .collect::<Vec<_>>(),
+            )),
+            ChangeColumn::CreatedAt => {
+                string_array(changes.iter().map(|row| Some(row.created_at.as_str())))
+            }
+            ChangeColumn::SnapshotContent => {
+                string_array(changes.iter().map(|row| row.snapshot_content.as_deref()))
+            }
+        })
+        .collect::<Vec<_>>();
+    record_batch_with_row_count(change_schema(projection), arrays, changes.len()).map_err(|error| {
+        DataFusionError::Execution(format!("failed to build lix_change batch: {error}"))
+    })
+}
+
+fn change_schema(projection: &[ChangeColumn]) -> SchemaRef {
+    Arc::new(Schema::new(
+        projection
+            .iter()
+            .map(|column| match column {
+                ChangeColumn::Id => Field::new("id", DataType::Utf8, false),
+                ChangeColumn::EntityId => Field::new("entity_id", DataType::Utf8, false),
+                ChangeColumn::SchemaKey => Field::new("schema_key", DataType::Utf8, false),
+                ChangeColumn::SchemaVersion => Field::new("schema_version", DataType::Utf8, false),
+                ChangeColumn::FileId => Field::new("file_id", DataType::Utf8, true),
+                ChangeColumn::Metadata => json_field("metadata", true),
+                ChangeColumn::CreatedAt => Field::new("created_at", DataType::Utf8, false),
+                ChangeColumn::SnapshotContent => json_field("snapshot_content", true),
+            })
+            .collect::<Vec<_>>(),
+    ))
+}
+
+fn string_array<'a>(values: impl Iterator<Item = Option<&'a str>>) -> ArrayRef {
+    Arc::new(StringArray::from(values.collect::<Vec<_>>())) as ArrayRef
+}
+
+fn datafusion_error_to_lix_error(error: DataFusionError) -> LixError {
+    super::error::datafusion_error_to_lix_error(error)
+}
+
+fn lix_error_to_datafusion_error(error: LixError) -> DataFusionError {
+    super::error::lix_error_to_datafusion_error(error)
+}

package/dist-engine-src/src/sql2/classify.rs

@@ -0,0 +1,147 @@
+use datafusion::sql::sqlparser::ast::{
+    FromTable, ObjectName, Query, SetExpr, Statement, TableFactor, TableObject, TableWithJoins,
+};
+use datafusion::sql::sqlparser::dialect::GenericDialect;
+use datafusion::sql::sqlparser::parser::Parser;
+
+use crate::LixError;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum SqlStatementKind {
+    Read,
+    Write,
+    Other,
+}
+
+pub(crate) fn classify_statement(sql: &str) -> Result<SqlStatementKind, LixError> {
+    let statements = parse_sql_statements(sql)?;
+    let [statement] = statements.as_slice() else {
+        return Ok(SqlStatementKind::Other);
+    };
+    Ok(classify_ast_statement(statement))
+}
+
+pub(crate) fn validate_supported_statement_ast(sql: &str) -> Result<(), LixError> {
+    let statements = parse_sql_statements(sql)?;
+    let [statement] = statements.as_slice() else {
+        return Err(unsupported_sql_error(
+            "Lix SQL only supports one statement per execute() call",
+        ));
+    };
+    validate_supported_ast_statement(statement)
+}
+
+pub(crate) fn dml_target_table_names(sql: &str) -> Result<Vec<String>, LixError> {
+    let statements = parse_sql_statements(sql)?;
+    let [statement] = statements.as_slice() else {
+        return Ok(Vec::new());
+    };
+    let mut targets = Vec::new();
+    collect_dml_target_table_names(statement, &mut targets);
+    Ok(targets)
+}
+
+fn parse_sql_statements(sql: &str) -> Result<Vec<Statement>, LixError> {
+    Parser::parse_sql(&GenericDialect {}, sql).map_err(|error| {
+        LixError::new(
+            LixError::CODE_PARSE_ERROR,
+            format!("sql2 SQL parse error: {error}"),
+        )
+    })
+}
+
+fn collect_dml_target_table_names(statement: &Statement, targets: &mut Vec<String>) {
+    match statement {
+        Statement::Insert(insert) => {
+            if let TableObject::TableName(name) = &insert.table {
+                if let Some(table_name) = object_name_table_part(name) {
+                    targets.push(table_name);
+                }
+            }
+        }
+        Statement::Update(update) => {
+            collect_table_with_joins_target(&update.table, targets);
+        }
+        Statement::Delete(delete) => {
+            let tables = match &delete.from {
+                FromTable::WithFromKeyword(tables) | FromTable::WithoutKeyword(tables) => tables,
+            };
+            for table in tables {
+                collect_table_with_joins_target(table, targets);
+            }
+        }
+        Statement::Explain { statement, .. } => {
+            collect_dml_target_table_names(statement.as_ref(), targets);
+        }
+        _ => {}
+    }
+}
+
+fn collect_table_with_joins_target(table: &TableWithJoins, targets: &mut Vec<String>) {
+    if let TableFactor::Table { name, .. } = &table.relation {
+        if let Some(table_name) = object_name_table_part(name) {
+            targets.push(table_name);
+        }
+    }
+}
+
+fn object_name_table_part(name: &ObjectName) -> Option<String> {
+    name.0
+        .last()
+        .and_then(|part| part.as_ident())
+        .map(|ident| ident.value.clone())
+}
+
+fn classify_ast_statement(statement: &Statement) -> SqlStatementKind {
+    match statement {
+        Statement::Insert(_) | Statement::Update(_) | Statement::Delete(_) => {
+            SqlStatementKind::Write
+        }
+        Statement::Query(_) => SqlStatementKind::Read,
+        Statement::Explain { statement, .. } => classify_ast_statement(statement.as_ref()),
+        _ => SqlStatementKind::Other,
+    }
+}
+
+fn validate_supported_ast_statement(statement: &Statement) -> Result<(), LixError> {
+    match statement {
+        Statement::Query(query) => validate_supported_query(query),
+        Statement::Insert(_) | Statement::Update(_) | Statement::Delete(_) => Ok(()),
+        Statement::Explain { statement, .. } => validate_supported_ast_statement(statement),
+        _ => Err(unsupported_sql_error(format!(
+            "SQL statement is not supported by Lix SQL: {statement}"
+        ))),
+    }
+}
+
+fn validate_supported_query(query: &Query) -> Result<(), LixError> {
+    if query.with.as_ref().is_some_and(|with| with.recursive) {
+        return Err(
+            unsupported_sql_error("recursive CTEs are not supported by Lix SQL").with_hint(
+                "Use explicit commit graph surfaces such as lix_commit, lix_commit_edge, and lix_state_history instead of WITH RECURSIVE.",
+            ),
+        );
+    }
+
+    if let Some(with) = &query.with {
+        for cte in &with.cte_tables {
+            validate_supported_query(&cte.query)?;
+        }
+    }
+    validate_supported_set_expr(&query.body)
+}
+
+fn validate_supported_set_expr(expr: &SetExpr) -> Result<(), LixError> {
+    match expr {
+        SetExpr::Query(query) => validate_supported_query(query),
+        SetExpr::SetOperation { left, right, .. } => {
+            validate_supported_set_expr(left)?;
+            validate_supported_set_expr(right)
+        }
+        _ => Ok(()),
+    }
+}
+
+fn unsupported_sql_error(message: impl Into<String>) -> LixError {
+    LixError::new(LixError::CODE_UNSUPPORTED_SQL, message)
+}