feat: Support for schema in connectors and datasources (#86)

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

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)

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

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")

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

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

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
 

src/ydata/sdk/datasources/_models/metadata/warnings.py

@@ -1,20 +1,13 @@
-from pydantic import BaseModel
-
+from ydata.sdk.common.model import BaseModel
 from ydata.sdk.datasources._models.metadata.warning_types import Level, WarningType
 
 
 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

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):

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

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")
 

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
 

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)