more update for the migration

ToucanToco · Jul 12, 2023 · 43cae1e · 43cae1e
1 parent 2d296ed
commit 43cae1e
Show file tree

Hide file tree

Showing 13 changed files with 51 additions and 32 deletions.
diff --git a/tests/test_connector.py b/tests/test_connector.py
@@ -30,8 +30,8 @@ class DataSource(ToucanDataSource):
 
 
 class DataConnector(ToucanConnector):
-    type = 'MyDB'
-    data_source_model: DataSource
+    type: str = 'MyDB'
+    data_source_model: DataSource = DataSource
     a_parameter: str = ''
 
     def _retrieve_data(self, data_source):
@@ -44,7 +44,7 @@ def test_missing_attributes():
     with pytest.raises(TypeError) as exc_info:
 
         class MissingDataConnector2(ToucanConnector):
-            type = 'MyDB'
+            type: str = 'MyDB'
 
             def _retrieve_data(self, data_source):
                 pass
@@ -54,7 +54,7 @@ def _retrieve_data(self, data_source):
 
 def test_no_get_df():
     class BadDataConnector(ToucanConnector):
-        type = 'MyDB'
+        type: str = 'MyDB'
         data_source_model = 'asd'
 
     with pytest.raises(TypeError):
@@ -75,7 +75,7 @@ def test_validate():
 
 def test_formated_engine_version():
     class DataConnector(ToucanConnector, VersionableEngineConnector):
-        type = 'MyDB'
+        type: str = 'MyDB'
         data_source_model: DataSource
 
         def get_engine_version(self) -> tuple:
@@ -100,7 +100,7 @@ def _retrieve_data(self, datasource):
 
 def test_get_df_with_permissions():
     class DataConnector(ToucanConnector):
-        type = 'MyDB'
+        type: str = 'MyDB'
         data_source_model: DataSource
 
         def _retrieve_data(self, datasource):
@@ -114,7 +114,7 @@ def _retrieve_data(self, datasource):
 
 def test_get_slice():
     class DataConnector(ToucanConnector):
-        type = 'MyDB'
+        type: str = 'MyDB'
         data_source_model = 'asd'
 
         def _retrieve_data(self, datasource):
@@ -145,7 +145,7 @@ def _retrieve_data(self, datasource):
 
 def test_explain():
     class DataConnector(ToucanConnector):
-        type = 'MyDB'
+        type: str = 'MyDB'
         data_source_model = 'asd'
 
         def _retrieve_data(self, datasource):
@@ -213,7 +213,7 @@ def test_get_cache_key_should_be_different_with_different_permissions():
 
 
 class UnreliableDataConnector(ToucanConnector):
-    type = 'MyUnreliableDB'
+    type: str = 'MyUnreliableDB'
     data_source_model: DataSource
 
     def _retrieve_data(self, data_source, logbook=[]):
@@ -232,7 +232,7 @@ def test_max_attempt_df():
 
 
 class CustomPolicyDataConnector(ToucanConnector):
-    type = 'MyUnreliableDB'
+    type: str = 'MyUnreliableDB'
     data_source_model: DataSource
 
     def _retrieve_data(self, data_source, logbook=[]):
@@ -254,7 +254,7 @@ def test_custom_max_attempt_df():
 
 
 class CustomRetryOnDataConnector(ToucanConnector):
-    type = 'MyUnreliableDB'
+    type: str = 'MyUnreliableDB'
     data_source_model: DataSource
     _retry_on = (ValueError,)
 
@@ -273,7 +273,7 @@ def test_custom_retry_on_df():
 
 
 class CustomNoRetryOnDataConnector(ToucanConnector):
-    type = 'MyUnreliableDB'
+    type: str = 'MyUnreliableDB'
     data_source_model: DataSource
 
     @property
@@ -345,7 +345,7 @@ def test_get_df_int_column(mocker):
     """The int column should be casted as str"""
 
     class DataConnector(ToucanConnector):
-        type = 'MyDB'
+        type: str = 'MyDB'
         data_source_model: DataSource
 
         def _retrieve_data(self, datasource):
@@ -357,7 +357,7 @@ def _retrieve_data(self, datasource):
 
 def test_default_implementation_of_discoverable_connector():
     class DataConnector(ToucanConnector, DiscoverableConnector):
-        type = 'MyDB'
+        type: str = 'MyDB'
         data_source_model: DataSource
 
         def _retrieve_data(self, datasource):

diff --git a/toucan_connectors/adobe_analytics/adobe_analytics_connector.py b/toucan_connectors/adobe_analytics/adobe_analytics_connector.py
@@ -3,8 +3,13 @@
 
 import pandas as pd
 from adobe_analytics import Client, ReportDefinition
+from pydantic import ConfigDict, Field
 
-from toucan_connectors.toucan_connector import ToucanConnector, ToucanDataSource
+from toucan_connectors.toucan_connector import (
+    PYDANTIC_VERSION_ONE,
+    ToucanConnector,
+    ToucanDataSource,
+)
 
 
 class Granularity(str, Enum):
@@ -49,12 +54,22 @@ class AdobeAnalyticsConnector(ToucanConnector):
     It provides a high-level interfaces for reporting queries (including Data Warehouse requests).
     """
 
-    data_source_model: AdobeAnalyticsDataSource
+    data_source_model: AdobeAnalyticsDataSource | None = None
 
     username: str
     password: str
     endpoint: str = Client.DEFAULT_ENDPOINT
 
+    # TODO[pydantic]: This is temporary, in the future we will only support V2
+    # and get rid of this condition + update the CI (link/test)
+    if PYDANTIC_VERSION_ONE:
+
+        class Config:
+            extra = 'allow'
+
+    else:
+        model_config = ConfigDict(extra='allow')
+
     def _retrieve_data(self, data_source: AdobeAnalyticsDataSource) -> pd.DataFrame:
         suites = Client(self.username, self.password, self.endpoint).suites()
         df = suites[data_source.suite_id].download(data_source.report_definition)

diff --git a/toucan_connectors/facebook_ads/facebook_ads_connector.py b/toucan_connectors/facebook_ads/facebook_ads_connector.py
@@ -81,7 +81,7 @@ class FacebookAdsConnector(ToucanConnector):
     auth_flow_id: Optional[str] = None
     _oauth_trigger = 'instance'
     data_source_model: FacebookAdsDataSource
-    oauth2_version = Field('1', **{'ui.hidden': True})
+    oauth2_version: str = Field('1', **{'ui.hidden': True})
     _oauth2_connector: OAuth2Connector = PrivateAttr()
 
     def __init__(self, **kwargs) -> None:

diff --git a/toucan_connectors/github/github_connector.py b/toucan_connectors/github/github_connector.py
@@ -82,7 +82,7 @@ class GithubDataSource(ToucanDataSource):
         description='Max Number of entities such as teams and repositories to extract',
     )
     _oauth_trigger = 'instance'
-    oauth2_version = Field('1', **{'ui.hidden': True})
+    oauth2_version: str = Field('1', **{'ui.hidden': True})
 
     @classmethod
     def get_form(cls, connector: 'GithubConnector', current_config, **kwargs):

diff --git a/toucan_connectors/google_adwords/google_adwords_connector.py b/toucan_connectors/google_adwords/google_adwords_connector.py
@@ -85,7 +85,7 @@ class GoogleAdwordsConnector(ToucanConnector):
     developer_token: str = None
     client_customer_id: str = None
     _oauth_trigger = 'instance'
-    oauth2_version = Field('1', **{'ui.hidden': True})
+    oauth2_version: str = Field('1', **{'ui.hidden': True})
     _oauth2_connector: OAuth2Connector = PrivateAttr()
 
     @staticmethod

diff --git a/toucan_connectors/google_sheets_2/google_sheets_2_connector.py b/toucan_connectors/google_sheets_2/google_sheets_2_connector.py
@@ -99,7 +99,7 @@ class GoogleSheets2Connector(ToucanConnector):
 
     _auth_flow = 'oauth2'
     _oauth_trigger = 'instance'
-    oauth2_version = Field('1', **{'ui.hidden': True})
+    oauth2_version: str = Field('1', **{'ui.hidden': True})
     auth_flow_id: Optional[str] = None
 
     # TODO: turn into a class property

diff --git a/toucan_connectors/hubspot/hubspot_connector.py b/toucan_connectors/hubspot/hubspot_connector.py
@@ -72,7 +72,7 @@ class HubspotConnector(ToucanConnector):
     _auth_flow = 'oauth2'
     auth_flow_id: Optional[str] = None
     _oauth_trigger = 'instance'
-    oauth2_version = Field('1', **{'ui.hidden': True})
+    oauth2_version: str = Field('1', **{'ui.hidden': True})
     data_source_model: HubspotDataSource
     _oauth2_connector: OAuth2Connector = PrivateAttr()
 

diff --git a/toucan_connectors/linkedinads/linkedinads_connector.py b/toucan_connectors/linkedinads/linkedinads_connector.py
@@ -105,7 +105,7 @@ class LinkedinadsConnector(ToucanConnector):
         description='You can provide a custom template that will be used for every HTTP request',
     )
     _oauth_trigger = 'instance'
-    oauth2_version = Field('1', **{'ui.hidden': True})
+    oauth2_version: str = Field('1', **{'ui.hidden': True})
     _oauth2_connector: OAuth2Connector = PrivateAttr()
 
     @staticmethod

diff --git a/toucan_connectors/one_drive/one_drive_connector.py b/toucan_connectors/one_drive/one_drive_connector.py
@@ -76,7 +76,7 @@ class OneDriveConnector(ToucanConnector):
 
     _auth_flow = 'oauth2'
     _oauth_trigger = 'connector'
-    oauth2_version = Field('1', **{'ui.hidden': True})
+    oauth2_version: str = Field('1', **{'ui.hidden': True})
     auth_flow_id: Optional[str] = None
 
     authorization_url: str = Field(None, **{'ui.hidden': True})

diff --git a/toucan_connectors/redshift/redshift_database_connector.py b/toucan_connectors/redshift/redshift_database_connector.py
@@ -15,6 +15,7 @@
 from toucan_connectors.redshift.utils import build_database_model_extraction_query, types_map
 from toucan_connectors.sql_query_helper import SqlQueryHelper
 from toucan_connectors.toucan_connector import (
+    PYDANTIC_VERSION_ONE,
     DataSlice,
     DiscoverableConnector,
     TableInfo,
@@ -163,7 +164,7 @@ def available_dbs(self) -> list[str]:
     def host_validator(cls, v):
         return re.sub(r'^https?://', '', v)
 
-    @root_validator
+    @root_validator(skip_on_failure=True)
     def check_requirements(cls, values):
         mode = values.get('authentication_method')
         if mode == AuthenticationMethod.DB_CREDENTIALS.value:

diff --git a/toucan_connectors/salesforce/salesforce_connector.py b/toucan_connectors/salesforce/salesforce_connector.py
@@ -59,7 +59,7 @@ class SalesforceConnector(ToucanConnector):
         description='Baseroute URL of the salesforces instance to query (without the trailing slash)',
     )
     _oauth_trigger = 'instance'
-    oauth2_version = Field('1', **{'ui.hidden': True})
+    oauth2_version: str = Field('1', **{'ui.hidden': True})
     _oauth2_connector: OAuth2Connector = PrivateAttr()
 
     @staticmethod

diff --git a/toucan_connectors/snowflake_oauth2/snowflake_oauth2_connector.py b/toucan_connectors/snowflake_oauth2/snowflake_oauth2_connector.py
@@ -81,7 +81,7 @@ class SnowflakeoAuth2Connector(ToucanConnector):
     auth_flow_id: str = Field(None, **{'ui.hidden': True})
     _auth_flow = 'oauth2'
     _oauth_trigger = 'connector'
-    oauth2_version = Field('1', **{'ui.hidden': True})
+    oauth2_version: str = Field('1', **{'ui.hidden': True})
     redirect_uri: str = Field(None, **{'ui.hidden': True})
     _oauth2_connector: OAuth2Connector = PrivateAttr()
 

diff --git a/toucan_connectors/toucan_connector.py b/toucan_connectors/toucan_connector.py
@@ -332,13 +332,16 @@ class Config:
             },
         )
 
+
     @classmethod
-    def __init_subclass__(cls):
-        try:
-            cls.data_source_model = cls.__fields__.pop('data_source_model').type_
-            cls.logger = logging.getLogger(cls.__name__)
-        except KeyError as e:
-            raise TypeError(f'{cls.__name__} has no {e} attribute.')
+    def __init_subclass__(cls, **kwargs: Any):
+        if PYDANTIC_VERSION_ONE:
+            try:
+                cls.data_source_model = cls.__fields__.pop('data_source_model').type_
+                cls.logger = logging.getLogger(cls.__name__)
+            except KeyError as e:
+                raise TypeError(f'{cls.__name__} has no {e} attribute.')
+
         if 'bearer_integration' in cls.__fields__:
             cls.bearer_integration = cls.__fields__['bearer_integration'].default