diff --git a/src/ydata/sdk/common/model.py b/src/ydata/sdk/common/model.py index 8c389515..7a2f6558 100644 --- a/src/ydata/sdk/common/model.py +++ b/src/ydata/sdk/common/model.py @@ -8,5 +8,6 @@ class BaseModel(PydanticBaseModel): All datamodel from YData SDK inherits from this class. """ class Config: - extra = Extra.ignore allow_population_by_field_name = True + extra = Extra.ignore + use_enum_values = True diff --git a/src/ydata/sdk/common/pydantic_utils.py b/src/ydata/sdk/common/pydantic_utils.py new file mode 100644 index 00000000..a19d3895 --- /dev/null +++ b/src/ydata/sdk/common/pydantic_utils.py @@ -0,0 +1,30 @@ +"""Alias generators for converting between different capitalization conventions.""" +import re + +__all__ = ('to_pascal', 'to_camel') + + +def to_pascal(snake: str) -> str: + """Convert a snake_case string to PascalCase. + + Args: + snake: The string to convert. + + Returns: + The PascalCase string. + """ + camel = snake.title() + return re.sub('([0-9A-Za-z])_(?=[0-9A-Z])', lambda m: m.group(1), camel) + + +def to_camel(snake: str) -> str: + """Convert a snake_case string to camelCase. + + Args: + snake: The string to convert. + + Returns: + The converted camelCase string. + """ + camel = to_pascal(snake) + return re.sub('(^_*[A-Z])', lambda m: m.group(1).lower(), camel) diff --git a/src/ydata/sdk/connectors/_models/connector_type.py b/src/ydata/sdk/connectors/_models/connector_type.py index 845c061e..c9c69c46 100644 --- a/src/ydata/sdk/connectors/_models/connector_type.py +++ b/src/ydata/sdk/connectors/_models/connector_type.py @@ -1,5 +1,8 @@ from enum import Enum +from typing import Union + +from ydata.sdk.common.exceptions import InvalidConnectorError class ConnectorType(Enum): @@ -19,3 +22,18 @@ class ConnectorType(Enum): """BigQuery connector""" SNOWFLAKE = "snowflake" """Snowflake connector""" + + @property + def is_rdbms(self) -> bool: + return self in [ConnectorType.AZURE_SQL, ConnectorType.MYSQL, ConnectorType.BIGQUERY, ConnectorType.SNOWFLAKE] + + @staticmethod + def _init_connector_type(connector_type: Union["ConnectorType", str]) -> "ConnectorType": + if isinstance(connector_type, str): + try: + connector_type = ConnectorType(connector_type) + except Exception: + c_list = ", ".join([c.value for c in ConnectorType]) + raise InvalidConnectorError( + f"ConnectorType '{connector_type}' does not exist.\nValid connector types are: {c_list}.") + return connector_type diff --git a/src/ydata/sdk/connectors/_models/rdbms_connector.py b/src/ydata/sdk/connectors/_models/rdbms_connector.py new file mode 100644 index 00000000..b8812337 --- /dev/null +++ b/src/ydata/sdk/connectors/_models/rdbms_connector.py @@ -0,0 +1,11 @@ +from typing import Optional + +from pydantic import Field + +from ydata.sdk.connectors._models.schema import Schema + +from .connector import Connector + + +class RDBMSConnector(Connector): + db_schema: Optional[Schema] = Field(None, alias="schema") diff --git a/src/ydata/sdk/connectors/_models/schema.py b/src/ydata/sdk/connectors/_models/schema.py new file mode 100644 index 00000000..e8696e91 --- /dev/null +++ b/src/ydata/sdk/connectors/_models/schema.py @@ -0,0 +1,43 @@ +from typing import List, Optional + +from pydantic import Field + +from ydata.sdk.common.model import BaseModel +from ydata.sdk.common.pydantic_utils import to_camel + + +class BaseConfig(BaseModel.Config): + alias_generator = to_camel + + +class TableColumn(BaseModel): + """Class to store the information of a Column table.""" + + name: str + variable_type: str # change this to the datatypes + primary_key: Optional[bool] + is_foreign_key: Optional[bool] + foreign_keys: list + nullable: bool + + Config = BaseConfig + + +class Table(BaseModel): + """Class to store the table columns information.""" + + name: str + columns: List[TableColumn] + primary_keys: List[TableColumn] + foreign_keys: List[TableColumn] + + Config = BaseConfig + + +class Schema(BaseModel): + """Class to store the database schema information.""" + + name: str + tables: Optional[List[Table]] = Field(None) + + Config = BaseConfig diff --git a/src/ydata/sdk/connectors/connector.py b/src/ydata/sdk/connectors/connector.py index 30a3e22f..ca7d58ee 100644 --- a/src/ydata/sdk/connectors/connector.py +++ b/src/ydata/sdk/connectors/connector.py @@ -6,13 +6,15 @@ from ydata.sdk.common.client import Client from ydata.sdk.common.client.utils import init_client from ydata.sdk.common.config import LOG_LEVEL -from ydata.sdk.common.exceptions import CredentialTypeError, InvalidConnectorError +from ydata.sdk.common.exceptions import CredentialTypeError from ydata.sdk.common.logger import create_logger from ydata.sdk.common.types import UID, Project from ydata.sdk.connectors._models.connector import Connector as mConnector from ydata.sdk.connectors._models.connector_list import ConnectorsList from ydata.sdk.connectors._models.connector_type import ConnectorType from ydata.sdk.connectors._models.credentials.credentials import Credentials +from ydata.sdk.connectors._models.rdbms_connector import RDBMSConnector as mRDBMSConnector +from ydata.sdk.connectors._models.schema import Schema from ydata.sdk.utils.model_mixin import ModelFactoryMixin @@ -33,11 +35,15 @@ class Connector(ModelFactoryMixin): type (ConnectorType): Type of the connector """ + _MODEL_CLASS = mConnector + + _model: Optional[mConnector] + def __init__( - self, connector_type: Union[ConnectorType, str] = None, credentials: Optional[Dict] = None, + self, connector_type: Union[ConnectorType, str, None] = None, credentials: Optional[Dict] = None, name: Optional[str] = None, project: Optional[Project] = None, client: Optional[Client] = None): self._init_common(client=client) - self._model: Optional[mConnector] = self._create_model( + self._model = _connector_type_to_model(ConnectorType._init_connector_type(connector_type))._create_model( connector_type, credentials, name, client=client) self._project = project @@ -61,7 +67,9 @@ def project(self) -> Project: @staticmethod @init_client - def get(uid: UID, project: Optional[Project] = None, client: Optional[Client] = None) -> "Connector": + def get( + uid: UID, project: Optional[Project] = None, client: Optional[Client] = None + ) -> Union["Connector", "RDBMSConnector"]: """Get an existing connector. Arguments: @@ -74,24 +82,19 @@ def get(uid: UID, project: Optional[Project] = None, client: Optional[Client] = """ response = client.get(f'/connector/{uid}', project=project) data = response.json() - connector = Connector._init_from_model_data(Connector, mConnector(**data)) + data_type = data["type"] + connector_class = _connector_type_to_model( + ConnectorType._init_connector_type(data_type)) + connector = connector_class._init_from_model_data( + connector_class._MODEL_CLASS(**data)) connector._project = project return connector @staticmethod - def _init_connector_type(connector_type: Union[ConnectorType, str]) -> ConnectorType: - if isinstance(connector_type, str): - try: - connector_type = ConnectorType(connector_type) - except Exception: - c_list = ", ".join([c.value for c in ConnectorType]) - raise InvalidConnectorError( - f"ConnectorType '{connector_type}' does not exist.\nValid connector types are: {c_list}.") - return connector_type - - @staticmethod - def _init_credentials(connector_type: ConnectorType, credentials: Union[str, Path, Dict, Credentials]) -> Credentials: + def _init_credentials( + connector_type: ConnectorType, credentials: Union[str, Path, Dict, Credentials] + ) -> Credentials: _credentials = None if isinstance(credentials, str): @@ -118,7 +121,7 @@ def _init_credentials(connector_type: ConnectorType, credentials: Union[str, Pat def create( connector_type: Union[ConnectorType, str], credentials: Union[str, Path, Dict, Credentials], name: Optional[str] = None, project: Optional[Project] = None, client: Optional[Client] = None - ) -> "Connector": + ) -> Union["Connector", "RDBMSConnector"]: """Create a new connector. Arguments: @@ -131,20 +134,22 @@ def create( Returns: New connector """ - model = Connector._create_model( + connector_type = ConnectorType._init_connector_type(connector_type) + connector_class = _connector_type_to_model(connector_type) + model = connector_class._create_model( connector_type=connector_type, credentials=credentials, name=name, project=project, client=client) - connector = ModelFactoryMixin._init_from_model_data( - Connector, model) + connector = connector_class._init_from_model_data(model) connector._project = project return connector @classmethod @init_client def _create_model( - cls, connector_type: Union[ConnectorType, str], credentials: Union[str, Path, Dict, Credentials], - name: Optional[str] = None, project: Optional[Project] = None, client: Optional[Client] = None) -> mConnector: + cls, connector_type: Union[ConnectorType, str], credentials: Union[str, Path, Dict, Credentials], + name: Optional[str] = None, project: Optional[Project] = None, client: Optional[Client] = None + ) -> Union[mConnector, mRDBMSConnector]: _name = name if name is not None else str(uuid4()) - _connector_type = Connector._init_connector_type(connector_type) + _connector_type = ConnectorType._init_connector_type(connector_type) _credentials = Connector._init_credentials(_connector_type, credentials) payload = { "type": _connector_type.value, @@ -154,7 +159,7 @@ def _create_model( response = client.post('/connector/', project=project, json=payload) data: list = response.json() - return mConnector(**data) + return cls._MODEL_CLASS(**data) @staticmethod @init_client @@ -174,3 +179,17 @@ def list(project: Optional[Project] = None, client: Optional[Client] = None) -> def __repr__(self): return self._model.__repr__() + + +class RDBMSConnector(Connector): + + _MODEL_CLASS = mRDBMSConnector + _model: Optional[mRDBMSConnector] + + @property + def schema(self) -> Optional[Schema]: + return self._model.db_schema + + +def _connector_type_to_model(connector_type: ConnectorType) -> Union[Connector, RDBMSConnector]: + return RDBMSConnector if connector_type.is_rdbms else Connector diff --git a/src/ydata/sdk/datasources/_models/metadata/metadata.py b/src/ydata/sdk/datasources/_models/metadata/metadata.py index 5cf34381..ea5b472f 100644 --- a/src/ydata/sdk/datasources/_models/metadata/metadata.py +++ b/src/ydata/sdk/datasources/_models/metadata/metadata.py @@ -1,7 +1,8 @@ from typing import List, Optional -from pydantic import BaseModel, Field +from pydantic import Field +from ydata.sdk.common.model import BaseModel from ydata.sdk.datasources._models.metadata.column import Column from ydata.sdk.datasources._models.metadata.warnings import MetadataWarning diff --git a/src/ydata/sdk/datasources/_models/metadata/warnings.py b/src/ydata/sdk/datasources/_models/metadata/warnings.py index 7eb39160..2968b99c 100644 --- a/src/ydata/sdk/datasources/_models/metadata/warnings.py +++ b/src/ydata/sdk/datasources/_models/metadata/warnings.py @@ -1,5 +1,4 @@ -from pydantic import BaseModel - +from ydata.sdk.common.model import BaseModel from ydata.sdk.datasources._models.metadata.warning_types import Level, WarningType @@ -7,14 +6,8 @@ class Details(BaseModel): level: Level value: str - class Config: - use_enum_values = True - class MetadataWarning(BaseModel): column: str details: Details type: WarningType - - class Config: - use_enum_values = True diff --git a/src/ydata/sdk/datasources/_models/status.py b/src/ydata/sdk/datasources/_models/status.py index 07739780..cc4fa4fa 100644 --- a/src/ydata/sdk/datasources/_models/status.py +++ b/src/ydata/sdk/datasources/_models/status.py @@ -1,6 +1,5 @@ -from pydantic import BaseModel - from ydata.core.enum import StringEnum +from ydata.sdk.common.model import BaseModel class ValidationState(StringEnum): diff --git a/src/ydata/sdk/datasources/datasource.py b/src/ydata/sdk/datasources/datasource.py index cb1f98f9..2d15777f 100644 --- a/src/ydata/sdk/datasources/datasource.py +++ b/src/ydata/sdk/datasources/datasource.py @@ -80,7 +80,7 @@ def status(self) -> Status: return Status.UNKNOWN @property - def metadata(self) -> Metadata: + def metadata(self) -> Optional[Metadata]: return self._model.metadata @staticmethod @@ -127,7 +127,7 @@ def get(uid: UID, project: Optional[Project] = None, client: Optional[Client] = datasource_type = CONNECTOR_TO_DATASOURCE.get( ConnectorType(data['connector']['type'])) model = DataSource._model_from_api(data, datasource_type) - datasource = ModelFactoryMixin._init_from_model_data(DataSource, model) + datasource = DataSource._init_from_model_data(model) datasource._project = project return datasource @@ -165,7 +165,7 @@ def _create( ) -> "DataSource": model = DataSource._create_model( connector, datasource_type, datatype, config, name, project, client) - datasource = ModelFactoryMixin._init_from_model_data(DataSource, model) + datasource = DataSource._init_from_model_data(model) if wait_for_metadata: datasource._model = DataSource._wait_for_metadata(datasource)._model diff --git a/src/ydata/sdk/synthesizers/_models/status.py b/src/ydata/sdk/synthesizers/_models/status.py index dd892bd5..673b7ef9 100644 --- a/src/ydata/sdk/synthesizers/_models/status.py +++ b/src/ydata/sdk/synthesizers/_models/status.py @@ -1,8 +1,9 @@ from typing import Generic, Optional, TypeVar -from pydantic import BaseModel, Field +from pydantic import Field from ydata.core.enum import StringEnum +from ydata.sdk.common.model import BaseModel T = TypeVar("T") diff --git a/src/ydata/sdk/synthesizers/_models/synthesizer.py b/src/ydata/sdk/synthesizers/_models/synthesizer.py index 7928c9a2..dd023195 100644 --- a/src/ydata/sdk/synthesizers/_models/synthesizer.py +++ b/src/ydata/sdk/synthesizers/_models/synthesizer.py @@ -1,6 +1,8 @@ from typing import Optional -from pydantic import BaseModel, Field +from pydantic import Field + +from ydata.sdk.common.model import BaseModel from .status import Status diff --git a/src/ydata/sdk/utils/model_mixin.py b/src/ydata/sdk/utils/model_mixin.py index 5598f4c0..c1eec025 100644 --- a/src/ydata/sdk/utils/model_mixin.py +++ b/src/ydata/sdk/utils/model_mixin.py @@ -1,4 +1,4 @@ -from typing import Any, Optional, Type +from typing import Any, Optional from ydata.sdk.common.client.client import Client from ydata.sdk.common.model import BaseModel @@ -8,8 +8,8 @@ class ModelFactoryMixin: """Mixin for class that implements an interface for an internal model. """ - @staticmethod - def _init_from_model_data(cls: Type, model: BaseModel, client: Optional[Client] = None) -> Any: + @classmethod + def _init_from_model_data(cls, model: BaseModel, client: Optional[Client] = None) -> Any: o = cls.__new__(cls) o._model = model o._init_common(client=client)