| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you 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. |
| from typing import Any, Dict, List, Optional |
| |
| from superset import app |
| from superset.models.core import Database |
| |
| custom_password_store = app.config["SQLALCHEMY_CUSTOM_PASSWORD_STORE"] |
| |
| |
| def get_foreign_keys_metadata( |
| database: Database, table_name: str, schema_name: Optional[str] |
| ) -> List[Dict[str, Any]]: |
| foreign_keys = database.get_foreign_keys(table_name, schema_name) |
| for fk in foreign_keys: |
| fk["column_names"] = fk.pop("constrained_columns") |
| fk["type"] = "fk" |
| return foreign_keys |
| |
| |
| def get_indexes_metadata( |
| database: Database, table_name: str, schema_name: Optional[str] |
| ) -> List[Dict[str, Any]]: |
| indexes = database.get_indexes(table_name, schema_name) |
| for idx in indexes: |
| idx["type"] = "index" |
| return indexes |
| |
| |
| def get_col_type(col: Dict[Any, Any]) -> str: |
| try: |
| dtype = f"{col['type']}" |
| except Exception: # pylint: disable=broad-except |
| # sqla.types.JSON __str__ has a bug, so using __class__. |
| dtype = col["type"].__class__.__name__ |
| return dtype |
| |
| |
| def get_table_metadata( |
| database: Database, table_name: str, schema_name: Optional[str] |
| ) -> Dict[str, Any]: |
| """ |
| Get table metadata information, including type, pk, fks. |
| This function raises SQLAlchemyError when a schema is not found. |
| |
| :param database: The database model |
| :param table_name: Table name |
| :param schema_name: schema name |
| :return: Dict table metadata ready for API response |
| """ |
| keys = [] |
| columns = database.get_columns(table_name, schema_name) |
| primary_key = database.get_pk_constraint(table_name, schema_name) |
| if primary_key and primary_key.get("constrained_columns"): |
| primary_key["column_names"] = primary_key.pop("constrained_columns") |
| primary_key["type"] = "pk" |
| keys += [primary_key] |
| foreign_keys = get_foreign_keys_metadata(database, table_name, schema_name) |
| indexes = get_indexes_metadata(database, table_name, schema_name) |
| keys += foreign_keys + indexes |
| payload_columns: List[Dict[str, Any]] = [] |
| table_comment = database.get_table_comment(table_name, schema_name) |
| for col in columns: |
| dtype = get_col_type(col) |
| payload_columns.append( |
| { |
| "name": col["name"], |
| "type": dtype.split("(")[0] if "(" in dtype else dtype, |
| "longType": dtype, |
| "keys": [k for k in keys if col["name"] in k["column_names"]], |
| "comment": col.get("comment"), |
| } |
| ) |
| return { |
| "name": table_name, |
| "columns": payload_columns, |
| "selectStar": database.select_star( |
| table_name, |
| schema=schema_name, |
| show_cols=True, |
| indent=True, |
| cols=columns, |
| latest_partition=True, |
| ), |
| "primaryKey": primary_key, |
| "foreignKeys": foreign_keys, |
| "indexes": keys, |
| "comment": table_comment, |
| } |