fix: table resolving logic related to pg_catalog

src/catalog/src/kvbackend/manager.rs

@@ -36,6 +36,7 @@ use futures_util::{StreamExt, TryStreamExt};
 use meta_client::client::MetaClient;
 use moka::sync::Cache;
 use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef};
+use session::context::{Channel, QueryContext};
 use snafu::prelude::*;
 use table::dist_table::DistTable;
 use table::table::numbers::{NumbersTable, NUMBERS_TABLE_NAME};
@@ -152,7 +153,11 @@ impl CatalogManager for KvBackendCatalogManager {
         Ok(keys)
     }
 
-    async fn schema_names(&self, catalog: &str) -> Result<Vec<String>> {
+    async fn schema_names(
+        &self,
+        catalog: &str,
+        query_ctx: Option<&QueryContext>,
+    ) -> Result<Vec<String>> {
         let stream = self
             .table_metadata_manager
             .schema_manager()
@@ -163,12 +168,17 @@ impl CatalogManager for KvBackendCatalogManager {
             .map_err(BoxedError::new)
             .context(ListSchemasSnafu { catalog })?;
 
-        keys.extend(self.system_catalog.schema_names());
+        keys.extend(self.system_catalog.schema_names(query_ctx));
 
         Ok(keys.into_iter().collect())
     }
 
-    async fn table_names(&self, catalog: &str, schema: &str) -> Result<Vec<String>> {
+    async fn table_names(
+        &self,
+        catalog: &str,
+        schema: &str,
+        query_ctx: Option<&QueryContext>,
+    ) -> Result<Vec<String>> {
         let stream = self
             .table_metadata_manager
             .table_name_manager()
@@ -181,7 +191,7 @@ impl CatalogManager for KvBackendCatalogManager {
             .into_iter()
             .map(|(k, _)| k)
             .collect::<Vec<_>>();
-        tables.extend_from_slice(&self.system_catalog.table_names(schema));
+        tables.extend_from_slice(&self.system_catalog.table_names(schema, query_ctx));
 
         Ok(tables.into_iter().collect())
     }
@@ -194,8 +204,13 @@ impl CatalogManager for KvBackendCatalogManager {
             .context(TableMetadataManagerSnafu)
     }
 
-    async fn schema_exists(&self, catalog: &str, schema: &str) -> Result<bool> {
-        if self.system_catalog.schema_exists(schema) {
+    async fn schema_exists(
+        &self,
+        catalog: &str,
+        schema: &str,
+        query_ctx: Option<&QueryContext>,
+    ) -> Result<bool> {
+        if self.system_catalog.schema_exists(schema, query_ctx) {
             return Ok(true);
         }
 
@@ -206,8 +221,14 @@ impl CatalogManager for KvBackendCatalogManager {
             .context(TableMetadataManagerSnafu)
     }
 
-    async fn table_exists(&self, catalog: &str, schema: &str, table: &str) -> Result<bool> {
-        if self.system_catalog.table_exists(schema, table) {
+    async fn table_exists(
+        &self,
+        catalog: &str,
+        schema: &str,
+        table: &str,
+        query_ctx: Option<&QueryContext>,
+    ) -> Result<bool> {
+        if self.system_catalog.table_exists(schema, table, query_ctx) {
             return Ok(true);
         }
 
@@ -225,34 +246,58 @@ impl CatalogManager for KvBackendCatalogManager {
         catalog_name: &str,
         schema_name: &str,
         table_name: &str,
+        query_ctx: Option<&QueryContext>,
     ) -> Result<Option<TableRef>> {
-        if let Some(table) = self
-            .system_catalog
-            .table(catalog_name, schema_name, table_name)
+        let channel = query_ctx.map_or(Channel::Unknown, |ctx| ctx.channel());
+        if let Some(table) =
+            self.system_catalog
+                .table(catalog_name, schema_name, table_name, query_ctx)
         {
             return Ok(Some(table));
         }
 
         let table_cache: TableCacheRef = self.cache_registry.get().context(CacheNotFoundSnafu {
             name: "table_cache",
         })?;
-
-        table_cache
+        if let Some(table) = table_cache
             .get_by_ref(&TableName {
                 catalog_name: catalog_name.to_string(),
                 schema_name: schema_name.to_string(),
                 table_name: table_name.to_string(),
             })
             .await
-            .context(GetTableCacheSnafu)
+            .context(GetTableCacheSnafu)?
+        {
+            return Ok(Some(table));
+        }
+
+        if channel == Channel::Postgres {
+            // falldown to pg_catalog
+            if let Some(table) =
+                self.system_catalog
+                    .table(catalog_name, PG_CATALOG_NAME, table_name, query_ctx)
+            {
+                return Ok(Some(table));
+            }
+        }
+
+        return Ok(None);
     }
 
-    fn tables<'a>(&'a self, catalog: &'a str, schema: &'a str) -> BoxStream<'a, Result<TableRef>> {
+    fn tables<'a>(
+        &'a self,
+        catalog: &'a str,
+        schema: &'a str,
+        query_ctx: Option<&'a QueryContext>,
+    ) -> BoxStream<'a, Result<TableRef>> {
         let sys_tables = try_stream!({
             // System tables
-            let sys_table_names = self.system_catalog.table_names(schema);
+            let sys_table_names = self.system_catalog.table_names(schema, query_ctx);
             for table_name in sys_table_names {
-                if let Some(table) = self.system_catalog.table(catalog, schema, &table_name) {
+                if let Some(table) =
+                    self.system_catalog
+                        .table(catalog, schema, &table_name, query_ctx)
+                {
                     yield table;
                 }
             }
@@ -320,42 +365,63 @@ struct SystemCatalog {
 }
 
 impl SystemCatalog {
-    // TODO(j0hn50n133): remove the duplicated hard-coded table names logic
-    fn schema_names(&self) -> Vec<String> {
-        vec![
-            INFORMATION_SCHEMA_NAME.to_string(),
-            PG_CATALOG_NAME.to_string(),
-        ]
+    fn schema_names(&self, query_ctx: Option<&QueryContext>) -> Vec<String> {
+        let channel = query_ctx.map_or(Channel::Unknown, |ctx| ctx.channel());
+        match channel {
+            // pg_catalog only visible under postgres protocol
+            Channel::Postgres => vec![
+                INFORMATION_SCHEMA_NAME.to_string(),
+                PG_CATALOG_NAME.to_string(),
+            ],
+            _ => {
+                vec![INFORMATION_SCHEMA_NAME.to_string()]
+            }
+        }
     }
 
-    fn table_names(&self, schema: &str) -> Vec<String> {
+    fn table_names(&self, schema: &str, query_ctx: Option<&QueryContext>) -> Vec<String> {
+        let channel = query_ctx.map_or(Channel::Unknown, |ctx| ctx.channel());
         match schema {
             INFORMATION_SCHEMA_NAME => self.information_schema_provider.table_names(),
-            PG_CATALOG_NAME => self.pg_catalog_provider.table_names(),
+            PG_CATALOG_NAME if channel == Channel::Postgres => {
+                self.pg_catalog_provider.table_names()
+            }
             DEFAULT_SCHEMA_NAME => {
                 vec![NUMBERS_TABLE_NAME.to_string()]
             }
             _ => vec![],
         }
     }
 
-    fn schema_exists(&self, schema: &str) -> bool {
-        schema == INFORMATION_SCHEMA_NAME || schema == PG_CATALOG_NAME
+    fn schema_exists(&self, schema: &str, query_ctx: Option<&QueryContext>) -> bool {
+        let channel = query_ctx.map_or(Channel::Unknown, |ctx| ctx.channel());
+        match channel {
+            Channel::Postgres => schema == PG_CATALOG_NAME || schema == INFORMATION_SCHEMA_NAME,
+            _ => schema == INFORMATION_SCHEMA_NAME,
+        }
     }
 
-    fn table_exists(&self, schema: &str, table: &str) -> bool {
+    fn table_exists(&self, schema: &str, table: &str, query_ctx: Option<&QueryContext>) -> bool {
+        let channel = query_ctx.map_or(Channel::Unknown, |ctx| ctx.channel());
         if schema == INFORMATION_SCHEMA_NAME {
             self.information_schema_provider.table(table).is_some()
         } else if schema == DEFAULT_SCHEMA_NAME {
             table == NUMBERS_TABLE_NAME
-        } else if schema == PG_CATALOG_NAME {
+        } else if schema == PG_CATALOG_NAME && channel == Channel::Postgres {
             self.pg_catalog_provider.table(table).is_some()
         } else {
             false
         }
     }
 
-    fn table(&self, catalog: &str, schema: &str, table_name: &str) -> Option<TableRef> {
+    fn table(
+        &self,
+        catalog: &str,
+        schema: &str,
+        table_name: &str,
+        query_ctx: Option<&QueryContext>,
+    ) -> Option<TableRef> {
+        let channel = query_ctx.map_or(Channel::Unknown, |ctx| ctx.channel());
         if schema == INFORMATION_SCHEMA_NAME {
             let information_schema_provider =
                 self.catalog_cache.get_with_by_ref(catalog, move || {
@@ -366,7 +432,7 @@ impl SystemCatalog {
                     ))
                 });
             information_schema_provider.table(table_name)
-        } else if schema == PG_CATALOG_NAME {
+        } else if schema == PG_CATALOG_NAME && channel == Channel::Postgres {
             if catalog == DEFAULT_CATALOG_NAME {
                 self.pg_catalog_provider.table(table_name)
             } else {

src/catalog/src/lib.rs

@@ -20,8 +20,10 @@ use std::fmt::{Debug, Formatter};
 use std::sync::Arc;
 
 use api::v1::CreateTableExpr;
+use common_catalog::consts::{INFORMATION_SCHEMA_NAME, PG_CATALOG_NAME};
 use futures::future::BoxFuture;
 use futures_util::stream::BoxStream;
+use session::context::QueryContext;
 use table::metadata::TableId;
 use table::TableRef;
 
@@ -44,26 +46,61 @@ pub trait CatalogManager: Send + Sync {
 
     async fn catalog_names(&self) -> Result<Vec<String>>;
 
-    async fn schema_names(&self, catalog: &str) -> Result<Vec<String>>;
+    async fn schema_names(
+        &self,
+        catalog: &str,
+        query_ctx: Option<&QueryContext>,
+    ) -> Result<Vec<String>>;
 
-    async fn table_names(&self, catalog: &str, schema: &str) -> Result<Vec<String>>;
+    async fn table_names(
+        &self,
+        catalog: &str,
+        schema: &str,
+        query_ctx: Option<&QueryContext>,
+    ) -> Result<Vec<String>>;
 
     async fn catalog_exists(&self, catalog: &str) -> Result<bool>;
 
-    async fn schema_exists(&self, catalog: &str, schema: &str) -> Result<bool>;
+    async fn schema_exists(
+        &self,
+        catalog: &str,
+        schema: &str,
+        query_ctx: Option<&QueryContext>,
+    ) -> Result<bool>;
 
-    async fn table_exists(&self, catalog: &str, schema: &str, table: &str) -> Result<bool>;
+    async fn table_exists(
+        &self,
+        catalog: &str,
+        schema: &str,
+        table: &str,
+        query_ctx: Option<&QueryContext>,
+    ) -> Result<bool>;
 
     /// Returns the table by catalog, schema and table name.
     async fn table(
         &self,
         catalog: &str,
         schema: &str,
         table_name: &str,
+        query_ctx: Option<&QueryContext>,
     ) -> Result<Option<TableRef>>;
 
     /// Returns all tables with a stream by catalog and schema.
-    fn tables<'a>(&'a self, catalog: &'a str, schema: &'a str) -> BoxStream<'a, Result<TableRef>>;
+    fn tables<'a>(
+        &'a self,
+        catalog: &'a str,
+        schema: &'a str,
+        query_ctx: Option<&'a QueryContext>,
+    ) -> BoxStream<'a, Result<TableRef>>;
+
+    /// Check if `schema` is a reserved schema name
+    fn is_reserved_schema_name(&self, schema: &str) -> bool {
+        // We have to check whether a schema name is reserved before create schema.
+        // We need this rather than use schema_exists directly because `pg_catalog` is
+        // only visible via postgres protocol. So if we don't check, a mysql client may
+        // create a schema named `pg_catalog` which is somehow malformed.
+        schema == INFORMATION_SCHEMA_NAME || schema == PG_CATALOG_NAME
+    }
 }
 
 pub type CatalogManagerRef = Arc<dyn CatalogManager>;