diff --git a/cloudquery/sdk/transformers/openapi.py b/cloudquery/sdk/transformers/openapi.py index 23af197..300acce 100644 --- a/cloudquery/sdk/transformers/openapi.py +++ b/cloudquery/sdk/transformers/openapi.py @@ -1,4 +1,4 @@ -from typing import Dict, List +from typing import Dict, List, Optional import pyarrow as pa from cloudquery.sdk.types import JSONType from cloudquery.sdk.schema import Column @@ -24,11 +24,24 @@ def oapi_type_to_arrow_type(field) -> pa.DataType: return pa.string() -def oapi_definition_to_columns(definition: Dict) -> List[Column]: +def get_column_by_name(columns: List[Column], name: str) -> Optional[Column]: + for column in columns: + if column.name == name: + return column + return None + + +def oapi_definition_to_columns(definition: Dict, override_columns=[]) -> List[Column]: columns = [] for key, value in definition["properties"].items(): column_type = oapi_type_to_arrow_type(value) - columns.append( - Column(name=key, type=column_type, description=value.get("description")) + column = Column( + name=key, type=column_type, description=value.get("description") ) + override_column = get_column_by_name(override_columns, key) + if override_column is not None: + column.type = override_column.type + column.primary_key = override_column.primary_key + column.unique = override_column.unique + columns.append(column) return columns