From 138846acca7f78b457c8e2be0809236b57ec2db0 Mon Sep 17 00:00:00 2001
From: Harshal Sheth <hsheth2@gmail.com>
Date: Tue, 30 Aug 2022 18:07:23 +0000
Subject: [PATCH] fix(ingest): add pydantic version upper bound (#5775)

---
 metadata-ingestion/setup.py                                  | 5 +++--
 metadata-ingestion/src/datahub/configuration/common.py       | 2 +-
 .../src/datahub/ingestion/source/azure/azure_common.py       | 5 +++++
 .../src/datahub/ingestion/source_config/sql/snowflake.py     | 2 +-
 4 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/metadata-ingestion/setup.py b/metadata-ingestion/setup.py
index 99e9192712..b9050f49c0 100644
--- a/metadata-ingestion/setup.py
+++ b/metadata-ingestion/setup.py
@@ -224,7 +224,7 @@ plugins: Dict[str, Set[str]] = {
     | bigquery_common
     | {"sqlalchemy-bigquery>=1.4.1", "sqllineage==1.3.5", "sqlparse"},
     "bigquery-usage": bigquery_common | usage_common | {"cachetools"},
-    "bigquery-beta": bigquery_common| {"sql_metadata"},
+    "bigquery-beta": bigquery_common | {"sql_metadata"},
     "clickhouse": sql_common | {"clickhouse-sqlalchemy==0.1.8"},
     "clickhouse-usage": sql_common
     | usage_common
@@ -350,7 +350,8 @@ base_dev_requirements = {
     "mypy>=0.950",
     # pydantic 1.8.2 is incompatible with mypy 0.910.
     # See https://github.com/samuelcolvin/pydantic/pull/3175#issuecomment-995382910.
-    "pydantic>=1.9.0",
+    # Restricting top version to <1.10 until we can fix our types.
+    "pydantic >=1.9.0, <1.10",
     "pytest>=6.2.2",
     "pytest-asyncio>=0.16.0",
     "pytest-cov>=2.8.1",
diff --git a/metadata-ingestion/src/datahub/configuration/common.py b/metadata-ingestion/src/datahub/configuration/common.py
index 1f5f033f9b..067370e4ba 100644
--- a/metadata-ingestion/src/datahub/configuration/common.py
+++ b/metadata-ingestion/src/datahub/configuration/common.py
@@ -105,7 +105,7 @@ class OauthConfiguration(ConfigModel):
     scopes: Optional[List[str]] = Field(
         description="scopes required to connect to snowflake"
     )
-    use_certificate: Optional[str] = Field(
+    use_certificate: Optional[bool] = Field(
         description="Do you want to use certificate and private key to authenticate using oauth",
         default=False,
     )
diff --git a/metadata-ingestion/src/datahub/ingestion/source/azure/azure_common.py b/metadata-ingestion/src/datahub/ingestion/source/azure/azure_common.py
index 491688704a..1a48725330 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/azure/azure_common.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/azure/azure_common.py
@@ -27,18 +27,23 @@ class AdlsSourceConfig(ConfigModel):
     )
     account_key: Optional[str] = Field(
         description="Azure storage account access key that can be used as a credential. **An account key, a SAS token or a client secret is required for authentication.**",
+        default=None,
     )
     sas_token: Optional[str] = Field(
         description="Azure storage account Shared Access Signature (SAS) token that can be used as a credential. **An account key, a SAS token or a client secret is required for authentication.**",
+        default=None,
     )
     client_secret: Optional[str] = Field(
         description="Azure client secret that can be used as a credential. **An account key, a SAS token or a client secret is required for authentication.**",
+        default=None,
     )
     client_id: Optional[str] = Field(
         description="Azure client (Application) ID required when a `client_secret` is used as a credential.",
+        default=None,
     )
     tenant_id: Optional[str] = Field(
         description="Azure tenant (Directory) ID required when a `client_secret` is used as a credential.",
+        default=None,
     )
 
     def get_abfss_url(self, folder_path: str = "") -> str:
diff --git a/metadata-ingestion/src/datahub/ingestion/source_config/sql/snowflake.py b/metadata-ingestion/src/datahub/ingestion/source_config/sql/snowflake.py
index 7953cb01bc..4018692ac8 100644
--- a/metadata-ingestion/src/datahub/ingestion/source_config/sql/snowflake.py
+++ b/metadata-ingestion/src/datahub/ingestion/source_config/sql/snowflake.py
@@ -332,7 +332,7 @@ class SnowflakeConfig(BaseSnowflakeConfig, SQLAlchemyConfig):
             self.oauth_config.authority_url,
             self.oauth_config.provider,
         )
-        if self.oauth_config.use_certificate is True:
+        if self.oauth_config.use_certificate:
             response = generator.get_token_with_certificate(
                 private_key_content=str(self.oauth_config.encoded_oauth_public_key),
                 public_key_content=str(self.oauth_config.encoded_oauth_private_key),