pipeline/manager/

table.rs

// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::sync::Arc;

use api::v1::value::ValueData;
use api::v1::{
    ColumnDataType, ColumnDef, ColumnSchema as PbColumnSchema, Row, RowInsertRequest,
    RowInsertRequests, Rows, SemanticType,
};
use arrow::array::{Array, AsArray};
use arrow::datatypes::TimestampNanosecondType;
use common_query::OutputData;
use common_recordbatch::util as record_util;
use common_telemetry::{debug, info};
use common_time::timestamp::{TimeUnit, Timestamp};
use datafusion::datasource::DefaultTableSource;
use datafusion::logical_expr::col;
use datafusion_common::TableReference;
use datafusion_expr::{DmlStatement, LogicalPlan};
use datatypes::timestamp::TimestampNanosecond;
use itertools::Itertools;
use operator::insert::InserterRef;
use operator::statement::StatementExecutorRef;
use query::QueryEngineRef;
use session::context::{QueryContextBuilder, QueryContextRef};
use snafu::{OptionExt, ResultExt, ensure};
use table::TableRef;
use table::metadata::TableInfo;
use table::table::adapter::DfTableProviderAdapter;

use crate::error::{
    BuildDfLogicalPlanSnafu, CastTypeSnafu, CollectRecordsSnafu, DataFrameSnafu, Error,
    ExecuteInternalStatementSnafu, InsertPipelineSnafu, InvalidPipelineVersionSnafu,
    MultiPipelineWithDiffSchemaSnafu, PipelineNotFoundSnafu, RecordBatchLenNotMatchSnafu, Result,
};
use crate::etl::{Content, Pipeline, parse_with_options};
use crate::manager::pipeline_cache::{PipelineCache, PipelineContent};
use crate::manager::{PipelineInfo, PipelineVersion};
use crate::metrics::METRIC_PIPELINE_TABLE_FIND_COUNT;
use crate::options::PipelineOptions;
use crate::util::prepare_dataframe_conditions;

pub(crate) const PIPELINE_TABLE_NAME: &str = "pipelines";
pub(crate) const PIPELINE_TABLE_PIPELINE_NAME_COLUMN_NAME: &str = "name";
const PIPELINE_TABLE_PIPELINE_SCHEMA_COLUMN_NAME: &str = "schema";
const PIPELINE_TABLE_PIPELINE_CONTENT_TYPE_COLUMN_NAME: &str = "content_type";
const PIPELINE_TABLE_PIPELINE_CONTENT_COLUMN_NAME: &str = "pipeline";
pub(crate) const PIPELINE_TABLE_CREATED_AT_COLUMN_NAME: &str = "created_at";
pub(crate) const EMPTY_SCHEMA_NAME: &str = "";

/// PipelineTable is a table that stores the pipeline schema and content.
/// Every catalog has its own pipeline table.
pub struct PipelineTable {
    inserter: InserterRef,
    statement_executor: StatementExecutorRef,
    table: TableRef,
    query_engine: QueryEngineRef,
    cache: PipelineCache,
    options: PipelineOptions,
}

impl PipelineTable {
    /// Create a new PipelineTable.
    pub fn new(
        inserter: InserterRef,
        statement_executor: StatementExecutorRef,
        table: TableRef,
        query_engine: QueryEngineRef,
        options: &PipelineOptions,
    ) -> Self {
        Self {
            inserter,
            statement_executor,
            table,
            query_engine,
            cache: PipelineCache::new(options.cache_ttl),
            options: options.clone(),
        }
    }

    /// Build the schema for the pipeline table.
    /// Returns the (time index, primary keys, column) definitions.
    pub fn build_pipeline_schema() -> (String, Vec<String>, Vec<ColumnDef>) {
        (
            PIPELINE_TABLE_CREATED_AT_COLUMN_NAME.to_string(),
            vec![
                PIPELINE_TABLE_PIPELINE_SCHEMA_COLUMN_NAME.to_string(),
                PIPELINE_TABLE_PIPELINE_NAME_COLUMN_NAME.to_string(),
                PIPELINE_TABLE_PIPELINE_CONTENT_TYPE_COLUMN_NAME.to_string(),
            ],
            vec![
                ColumnDef {
                    name: PIPELINE_TABLE_PIPELINE_NAME_COLUMN_NAME.to_string(),
                    data_type: ColumnDataType::String as i32,
                    is_nullable: false,
                    default_constraint: vec![],
                    semantic_type: SemanticType::Tag as i32,
                    comment: "".to_string(),
                    datatype_extension: None,
                    options: None,
                },
                ColumnDef {
                    name: PIPELINE_TABLE_PIPELINE_SCHEMA_COLUMN_NAME.to_string(),
                    data_type: ColumnDataType::String as i32,
                    is_nullable: false,
                    default_constraint: vec![],
                    semantic_type: SemanticType::Tag as i32,
                    comment: "".to_string(),
                    datatype_extension: None,
                    options: None,
                },
                ColumnDef {
                    name: PIPELINE_TABLE_PIPELINE_CONTENT_TYPE_COLUMN_NAME.to_string(),
                    data_type: ColumnDataType::String as i32,
                    is_nullable: false,
                    default_constraint: vec![],
                    semantic_type: SemanticType::Tag as i32,
                    comment: "".to_string(),
                    datatype_extension: None,
                    options: None,
                },
                ColumnDef {
                    name: PIPELINE_TABLE_PIPELINE_CONTENT_COLUMN_NAME.to_string(),
                    data_type: ColumnDataType::String as i32,
                    is_nullable: false,
                    default_constraint: vec![],
                    semantic_type: SemanticType::Field as i32,
                    comment: "".to_string(),
                    datatype_extension: None,
                    options: None,
                },
                ColumnDef {
                    name: PIPELINE_TABLE_CREATED_AT_COLUMN_NAME.to_string(),
                    data_type: ColumnDataType::TimestampNanosecond as i32,
                    is_nullable: false,
                    default_constraint: vec![],
                    semantic_type: SemanticType::Timestamp as i32,
                    comment: "".to_string(),
                    datatype_extension: None,
                    options: None,
                },
            ],
        )
    }

    /// Build the column schemas for inserting a row into the pipeline table.
    fn build_insert_column_schemas() -> Vec<PbColumnSchema> {
        vec![
            PbColumnSchema {
                column_name: PIPELINE_TABLE_PIPELINE_NAME_COLUMN_NAME.to_string(),
                datatype: ColumnDataType::String.into(),
                semantic_type: SemanticType::Tag.into(),
                ..Default::default()
            },
            PbColumnSchema {
                column_name: PIPELINE_TABLE_PIPELINE_SCHEMA_COLUMN_NAME.to_string(),
                datatype: ColumnDataType::String.into(),
                semantic_type: SemanticType::Tag.into(),
                ..Default::default()
            },
            PbColumnSchema {
                column_name: PIPELINE_TABLE_PIPELINE_CONTENT_TYPE_COLUMN_NAME.to_string(),
                datatype: ColumnDataType::String.into(),
                semantic_type: SemanticType::Tag.into(),
                ..Default::default()
            },
            PbColumnSchema {
                column_name: PIPELINE_TABLE_PIPELINE_CONTENT_COLUMN_NAME.to_string(),
                datatype: ColumnDataType::String.into(),
                semantic_type: SemanticType::Field.into(),
                ..Default::default()
            },
            PbColumnSchema {
                column_name: PIPELINE_TABLE_CREATED_AT_COLUMN_NAME.to_string(),
                datatype: ColumnDataType::TimestampNanosecond.into(),
                semantic_type: SemanticType::Timestamp.into(),
                ..Default::default()
            },
        ]
    }

    fn query_ctx(table_info: &TableInfo) -> QueryContextRef {
        QueryContextBuilder::default()
            .current_catalog(table_info.catalog_name.clone())
            .current_schema(table_info.schema_name.clone())
            .build()
            .into()
    }

    /// Compile a pipeline from a string.
    pub fn compile_pipeline(&self, pipeline: &str) -> Result<Pipeline> {
        let yaml_content = Content::Yaml(pipeline);
        parse_with_options(&yaml_content, &self.options)
    }

    /// Insert a pipeline into the pipeline table.
    async fn insert_pipeline_to_pipeline_table(
        &self,
        name: &str,
        content_type: &str,
        pipeline: &str,
    ) -> Result<Timestamp> {
        let now = Timestamp::current_time(TimeUnit::Nanosecond);

        let table_info = self.table.table_info();

        let insert = RowInsertRequest {
            table_name: PIPELINE_TABLE_NAME.to_string(),
            rows: Some(Rows {
                schema: Self::build_insert_column_schemas(),
                rows: vec![Row {
                    values: vec![
                        ValueData::StringValue(name.to_string()).into(),
                        ValueData::StringValue(EMPTY_SCHEMA_NAME.to_string()).into(),
                        ValueData::StringValue(content_type.to_string()).into(),
                        ValueData::StringValue(pipeline.to_string()).into(),
                        ValueData::TimestampNanosecondValue(now.value()).into(),
                    ],
                }],
            }),
        };

        let requests = RowInsertRequests {
            inserts: vec![insert],
        };

        let output = self
            .inserter
            .handle_row_inserts(
                requests,
                Self::query_ctx(&table_info),
                &self.statement_executor,
                false,
                false,
            )
            .await
            .context(InsertPipelineSnafu)?;

        info!(
            "Insert pipeline success, name: {:?}, table: {:?}, output: {:?}",
            name,
            table_info.full_table_name(),
            output
        );

        Ok(now)
    }

    /// Get a pipeline by name.
    /// If the pipeline is not in the cache, it will be get from table and compiled and inserted into the cache.
    pub async fn get_pipeline(
        &self,
        schema: &str,
        name: &str,
        input_version: PipelineVersion,
    ) -> Result<Arc<Pipeline>> {
        self.cache
            .get_pipeline_with(schema, name, input_version, async {
                let pipeline_content = self.get_pipeline_str(schema, name, input_version).await?;
                Ok(Arc::new(self.compile_pipeline(&pipeline_content.content)?))
            })
            .await
    }

    /// Get a original pipeline by name.
    /// If the pipeline is not in the cache, it will be get from table and compiled and inserted into the cache.
    pub async fn get_pipeline_str(
        &self,
        schema: &str,
        name: &str,
        input_version: PipelineVersion,
    ) -> Result<PipelineContent> {
        self.cache
            .get_pipeline_str_with(schema, name, input_version, async {
                self.load_pipeline_str(schema, name, input_version).await
            })
            .await
    }

    async fn load_pipeline_str(
        &self,
        schema: &str,
        name: &str,
        input_version: PipelineVersion,
    ) -> Result<PipelineContent> {
        let mut pipeline_vec;
        match self.find_pipeline(name, input_version).await {
            Ok(p) => {
                METRIC_PIPELINE_TABLE_FIND_COUNT
                    .with_label_values(&["true"])
                    .inc();
                pipeline_vec = p;
            }
            Err(e) => {
                match e {
                    Error::CollectRecords { .. } => {
                        // if collect records failed, it means the pipeline table is temporary invalid
                        // we should use failover cache
                        METRIC_PIPELINE_TABLE_FIND_COUNT
                            .with_label_values(&["false"])
                            .inc();
                        return self
                            .cache
                            .get_failover_cache(schema, name, input_version)
                            .await?
                            .context(PipelineNotFoundSnafu {
                                name,
                                version: input_version,
                            });
                    }
                    _ => {
                        // if other error, we should return it
                        return Err(e);
                    }
                }
            }
        };
        ensure!(
            !pipeline_vec.is_empty(),
            PipelineNotFoundSnafu {
                name,
                version: input_version
            }
        );

        // if the result is exact one, use it
        if pipeline_vec.len() == 1 {
            let pipeline_content = pipeline_vec.remove(0);

            self.cache
                .insert_failover_cache(pipeline_content.clone(), input_version.is_none())
                .await;
            return Ok(pipeline_content);
        }

        // check if there's empty schema pipeline
        // if there isn't, check current schema
        let pipeline = pipeline_vec
            .iter()
            .position(|v| v.schema == EMPTY_SCHEMA_NAME)
            .or_else(|| pipeline_vec.iter().position(|v| v.schema == schema))
            .map(|idx| pipeline_vec.remove(idx));

        // multiple pipeline with no empty or current schema
        // throw an error
        let pipeline_content = pipeline.with_context(|| MultiPipelineWithDiffSchemaSnafu {
            name: name.to_string(),
            current_schema: schema.to_string(),
            schemas: pipeline_vec.iter().map(|v| v.schema.clone()).join(","),
        })?;

        self.cache
            .insert_failover_cache(pipeline_content.clone(), input_version.is_none())
            .await;
        Ok(pipeline_content)
    }

    /// Insert a pipeline into the pipeline table and compile it.
    /// Newly created pipelines will be saved under empty schema.
    pub async fn insert_and_compile(
        &self,
        name: &str,
        content_type: &str,
        pipeline: &str,
    ) -> Result<PipelineInfo> {
        let compiled_pipeline = Arc::new(self.compile_pipeline(pipeline)?);
        // we will use the version in the future
        let version = self
            .insert_pipeline_to_pipeline_table(name, content_type, pipeline)
            .await?;

        self.cache
            .on_pipeline_created(PipelineContent {
                name: name.to_string(),
                content: pipeline.to_string(),
                version: TimestampNanosecond(version),
                schema: EMPTY_SCHEMA_NAME.to_string(),
            })
            .await;

        Ok((version, compiled_pipeline))
    }

    pub async fn delete_pipeline(
        &self,
        name: &str,
        version: PipelineVersion,
    ) -> Result<Option<()>> {
        // 0. version is ensured at the http api level not None
        ensure!(
            version.is_some(),
            InvalidPipelineVersionSnafu { version: "None" }
        );

        // 1. check pipeline exist in catalog
        let pipeline = self.find_pipeline(name, version).await?;
        if pipeline.is_empty() {
            return Ok(None);
        }

        // 2. prepare dataframe
        let dataframe = self
            .query_engine
            .read_table(self.table.clone())
            .context(DataFrameSnafu)?;

        let dataframe = dataframe
            .filter(prepare_dataframe_conditions(name, version))
            .context(BuildDfLogicalPlanSnafu)?;

        // 3. prepare dml stmt
        let table_info = self.table.table_info();
        let table_name = TableReference::full(
            table_info.catalog_name.clone(),
            table_info.schema_name.clone(),
            table_info.name.clone(),
        );

        let table_provider = Arc::new(DfTableProviderAdapter::new(self.table.clone()));
        let table_source = Arc::new(DefaultTableSource::new(table_provider));

        // create dml stmt
        let stmt = DmlStatement::new(
            table_name,
            table_source,
            datafusion_expr::WriteOp::Delete,
            Arc::new(dataframe.into_parts().1),
        );

        let plan = LogicalPlan::Dml(stmt);

        // 4. execute dml stmt
        let output = self
            .query_engine
            .execute(plan, Self::query_ctx(&table_info))
            .await
            .context(ExecuteInternalStatementSnafu)?;

        info!(
            "Delete pipeline success, name: {:?}, version: {:?}, table: {:?}, output: {:?}",
            name,
            version,
            table_info.full_table_name(),
            output
        );

        self.cache.invalidate(name, version).await;

        Ok(Some(()))
    }

    // find all pipelines with name and version
    // cloud be multiple with different schema
    // return format: (pipeline content, schema, created_at)
    async fn find_pipeline(
        &self,
        name: &str,
        version: PipelineVersion,
    ) -> Result<Vec<PipelineContent>> {
        // 1. prepare dataframe
        let dataframe = self
            .query_engine
            .read_table(self.table.clone())
            .context(DataFrameSnafu)?;

        // select all pipelines with name and version
        let dataframe = dataframe
            .filter(prepare_dataframe_conditions(name, version))
            .context(BuildDfLogicalPlanSnafu)?
            .select_columns(&[
                PIPELINE_TABLE_PIPELINE_CONTENT_COLUMN_NAME,
                PIPELINE_TABLE_PIPELINE_SCHEMA_COLUMN_NAME,
                PIPELINE_TABLE_CREATED_AT_COLUMN_NAME,
            ])
            .context(BuildDfLogicalPlanSnafu)?
            .sort(vec![
                col(PIPELINE_TABLE_CREATED_AT_COLUMN_NAME).sort(false, true),
            ])
            .context(BuildDfLogicalPlanSnafu)?;

        let plan = dataframe.into_parts().1;

        let table_info = self.table.table_info();

        debug!("find_pipeline_by_name: plan: {:?}", plan);

        // 2. execute plan
        let output = self
            .query_engine
            .execute(plan, Self::query_ctx(&table_info))
            .await
            .context(ExecuteInternalStatementSnafu)?;
        let output = output
            .map_dictionary_to_values()
            .context(CollectRecordsSnafu)?;
        let stream = match output.data {
            OutputData::Stream(stream) => stream,
            OutputData::RecordBatches(record_batches) => record_batches.as_stream(),
            _ => unreachable!(),
        };

        // 3. construct result
        let records = record_util::collect(stream)
            .await
            .context(CollectRecordsSnafu)?;

        if records.is_empty() {
            return Ok(vec![]);
        }

        ensure!(
            !records.is_empty() && records.iter().all(|r| r.num_columns() == 3),
            PipelineNotFoundSnafu { name, version }
        );

        let mut re = Vec::with_capacity(records.len());
        for r in records {
            let pipeline_content_column = r.column(0);
            let pipeline_content = pipeline_content_column
                .as_string_opt::<i32>()
                .with_context(|| CastTypeSnafu {
                    msg: format!(
                        "can't downcast {:?} array into string vector",
                        pipeline_content_column.data_type()
                    ),
                })?;

            let pipeline_schema_column = r.column(1);
            let pipeline_schema =
                pipeline_schema_column
                    .as_string_opt::<i32>()
                    .with_context(|| CastTypeSnafu {
                        msg: format!(
                            "expecting pipeline schema column of type string, actual: {}",
                            pipeline_schema_column.data_type()
                        ),
                    })?;

            let pipeline_created_at_column = r.column(2);
            let pipeline_created_at = pipeline_created_at_column
                .as_primitive_opt::<TimestampNanosecondType>()
                .with_context(|| CastTypeSnafu {
                    msg: format!(
                        "can't downcast {:?} array into scalar vector",
                        pipeline_created_at_column.data_type()
                    ),
                })?;

            debug!(
                "find_pipeline_by_name: pipeline_content: {:?}, pipeline_schema: {:?}, pipeline_created_at: {:?}",
                pipeline_content, pipeline_schema, pipeline_created_at
            );

            ensure!(
                pipeline_content.len() == pipeline_schema.len()
                    && pipeline_schema.len() == pipeline_created_at.len(),
                RecordBatchLenNotMatchSnafu
            );

            let len = pipeline_content.len();
            for i in 0..len {
                re.push(PipelineContent {
                    name: name.to_string(),
                    content: pipeline_content.value(i).to_string(),
                    version: TimestampNanosecond::new(pipeline_created_at.value(i)),
                    schema: pipeline_schema.value(i).to_string(),
                });
            }
        }

        Ok(re)
    }
}