From cb72a22b59c036d20181bb299840344a0eedb0e4 Mon Sep 17 00:00:00 2001 From: Pere Miquel Brull Date: Fri, 7 Jun 2024 04:36:17 +0200 Subject: [PATCH] Fix - e2e tests for pydantic V2 (#16551) * Fix - e2e tests for pydantic V2 * add correct default * add correct default * revert datetime aware * revert datetime aware * revert datetime aware * revert datetime aware * revert datetime aware * revert datetime aware * revert datetime aware * revert datetime aware * fix apis * format --- ingestion/operators/docker/exit_handler.py | 2 +- ingestion/operators/docker/run_automation.py | 2 +- .../processor/reports/data_processor.py | 4 +- .../pandas/pandas_test_suite_interface.py | 4 +- .../sqlalchemy/sqa_test_suite_interface.py | 4 +- .../src/metadata/great_expectations/action.py | 6 +- .../ingestion/bulksink/metadata_usage.py | 9 ++- .../ingestion/connections/test_connections.py | 8 +-- .../metadata/ingestion/ometa/auth_provider.py | 4 +- .../ingestion/ometa/mixins/tests_mixin.py | 4 +- .../source/dashboard/looker/metadata.py | 4 +- .../source/dashboard/metabase/client.py | 2 +- .../source/dashboard/tableau/connection.py | 4 +- .../source/database/bigquery/connection.py | 4 +- .../source/database/bigquery/metadata.py | 6 +- .../source/database/databricks/lineage.py | 4 +- .../source/database/databricks/metadata.py | 4 +- .../source/database/databricks/usage.py | 4 +- .../incremental_metadata_extraction.py | 2 +- .../source/database/life_cycle_query_mixin.py | 9 ++- .../source/database/postgres/lineage.py | 4 +- .../source/database/postgres/usage.py | 4 +- .../ingestion/source/database/sample_data.py | 28 +++----- .../ingestion/source/database/sample_usage.py | 4 +- .../ingestion/source/database/sas/metadata.py | 10 ++- .../database/stored_procedures_mixin.py | 11 ++- .../ingestion/source/database/usage_source.py | 17 ++--- .../source/pipeline/kafkaconnect/metadata.py | 4 +- .../ingestion/source/storage/s3/metadata.py | 6 +- .../interface/nosql/profiler_interface.py | 6 +- .../interface/pandas/profiler_interface.py | 6 +- .../sqlalchemy/profiler_interface.py | 6 +- .../src/metadata/profiler/processor/core.py | 4 +- .../profiler/processor/sample_data_handler.py | 4 +- ingestion/src/metadata/utils/tag_utils.py | 8 +-- ingestion/src/metadata/workflow/base.py | 4 +- .../workflow/workflow_status_mixin.py | 4 +- ingestion/tests/cli_e2e/base/test_cli.py | 14 +++- .../cli_e2e/dashboard/tableau/tableau.yaml | 2 +- .../cli_e2e/database/bigquery/bigquery.yaml | 2 +- .../bigquery_multiple_project.yaml | 2 +- .../data_insight/producer/test_producers.py | 6 +- .../test_data_insight_workflow.py | 2 +- .../data_insight/test_web_analytic_events.py | 2 +- .../ometa/test_ometa_lineage_api.py | 2 +- .../integration/ometa/test_ometa_table_api.py | 14 ++-- .../system/test_bigquery_system_metrics.py | 4 +- .../system/test_redshift_system_metrics.py | 4 +- .../system/test_snowflake_system_metrics.py | 4 +- .../test_suite/test_e2e_workflow.py | 2 +- ...st_registry_names_match_test_definition.py | 2 +- .../integration/test_suite/test_workflow.py | 2 +- .../pandas/test_profiler_interface.py | 4 +- .../sqlalchemy/bigquery/test_map_struct.py | 2 +- .../sqlalchemy/test_sqa_profiler_interface.py | 4 +- .../tests/unit/test_databricks_lineage.py | 12 ++-- ingestion/tests/unit/test_dbt.py | 2 +- .../tests/unit/test_handle_partitions.py | 2 +- .../tests/unit/test_ometa_to_dataframe.py | 4 +- .../tests/unit/test_pgspider_lineage_unit.py | 2 +- ingestion/tests/unit/test_pydantic_v2.py | 4 +- ingestion/tests/unit/test_source_parsing.py | 70 +++++++++---------- ingestion/tests/unit/test_usage_filter.py | 2 +- ingestion/tests/unit/test_usage_log.py | 2 +- .../topology/dashboard/test_domodashboard.py | 4 +- .../unit/topology/dashboard/test_looker.py | 2 +- .../unit/topology/dashboard/test_metabase.py | 2 +- .../unit/topology/dashboard/test_qlikcloud.py | 4 +- .../unit/topology/dashboard/test_qliksense.py | 4 +- .../topology/dashboard/test_quicksight.py | 2 +- .../unit/topology/dashboard/test_superset.py | 4 +- .../unit/topology/dashboard/test_tableau.py | 2 +- .../unit/topology/database/test_bigquery.py | 4 +- .../unit/topology/database/test_couchbase.py | 2 +- .../unit/topology/database/test_databricks.py | 2 +- .../unit/topology/database/test_datalake.py | 10 +-- .../unit/topology/database/test_deltalake.py | 2 +- .../topology/database/test_domodatabase.py | 4 +- .../unit/topology/database/test_doris.py | 2 +- .../tests/unit/topology/database/test_glue.py | 2 +- .../unit/topology/database/test_greenplum.py | 2 +- .../tests/unit/topology/database/test_hive.py | 2 +- .../unit/topology/database/test_iceberg.py | 6 +- .../unit/topology/database/test_mongodb.py | 2 +- .../unit/topology/database/test_mssql.py | 2 +- .../unit/topology/database/test_mysql.py | 2 +- .../unit/topology/database/test_oracle.py | 4 +- .../unit/topology/database/test_postgres.py | 4 +- .../unit/topology/database/test_redshift.py | 2 +- .../unit/topology/database/test_salesforce.py | 2 +- .../tests/unit/topology/database/test_sas.py | 4 +- .../unit/topology/database/test_snowflake.py | 4 +- .../unit/topology/metadata/test_amundsen.py | 2 +- .../unit/topology/metadata/test_atlas.py | 4 +- .../unit/topology/pipeline/test_airbyte.py | 2 +- .../unit/topology/pipeline/test_airflow.py | 4 +- .../unit/topology/pipeline/test_dagster.py | 2 +- .../pipeline/test_databricks_pipeline.py | 2 +- .../topology/pipeline/test_domopipeline.py | 8 ++- .../unit/topology/pipeline/test_fivetran.py | 2 +- .../tests/unit/topology/pipeline/test_nifi.py | 2 +- .../topology/pipeline/test_openlineage.py | 2 +- .../unit/topology/pipeline/test_spline.py | 2 +- .../topology/search/test_elasticsearch.py | 2 +- .../unit/topology/storage/test_storage.py | 6 +- .../operations/deploy.py | 3 +- .../workflows/ingestion/application.py | 3 +- .../workflows/ingestion/common.py | 19 ++--- .../workflows/ingestion/data_insight.py | 3 +- .../workflows/ingestion/profiler.py | 3 +- .../workflows/ingestion/test_suite.py | 3 +- .../workflows/ingestion/usage.py | 3 +- .../test_workflow_creation.py | 26 ++++--- .../metadata/openMetadataConnection.json | 2 +- scripts/datamodel_generation.py | 18 +++++ 115 files changed, 316 insertions(+), 289 deletions(-) diff --git a/ingestion/operators/docker/exit_handler.py b/ingestion/operators/docker/exit_handler.py index 5d2febfb9bcc..9dbe7369b554 100644 --- a/ingestion/operators/docker/exit_handler.py +++ b/ingestion/operators/docker/exit_handler.py @@ -74,7 +74,7 @@ def main(): raw_workflow_config = yaml.safe_load(config) raw_workflow_config["pipelineRunId"] = pipeline_run_id - workflow_config = OpenMetadataWorkflowConfig.parse_obj(raw_workflow_config) + workflow_config = OpenMetadataWorkflowConfig.model_validate(raw_workflow_config) metadata = OpenMetadata( config=workflow_config.workflowConfig.openMetadataServerConfig ) diff --git a/ingestion/operators/docker/run_automation.py b/ingestion/operators/docker/run_automation.py index 7c9d8b72476b..513b4b5857c7 100644 --- a/ingestion/operators/docker/run_automation.py +++ b/ingestion/operators/docker/run_automation.py @@ -65,7 +65,7 @@ def main(): set_loggers_level(logging.INFO) automation_workflow_dict = yaml.safe_load(config) - automation_workflow = AutomationWorkflow.parse_obj(automation_workflow_dict) + automation_workflow = AutomationWorkflow.model_validate(automation_workflow_dict) execute(automation_workflow) diff --git a/ingestion/src/metadata/data_insight/processor/reports/data_processor.py b/ingestion/src/metadata/data_insight/processor/reports/data_processor.py index b4b3156b1df6..1b5a6675523d 100644 --- a/ingestion/src/metadata/data_insight/processor/reports/data_processor.py +++ b/ingestion/src/metadata/data_insight/processor/reports/data_processor.py @@ -16,7 +16,7 @@ from __future__ import annotations import abc -from datetime import datetime, timezone +from datetime import datetime from typing import Callable, Iterable, Optional from metadata.generated.schema.analytics.reportData import ReportData @@ -43,7 +43,7 @@ def __init_subclass__(cls, *args, **kwargs) -> None: def __init__(self, metadata: OpenMetadata): self.metadata = metadata - self.timestamp = Timestamp(int(datetime.now(timezone.utc).timestamp() * 1000)) + self.timestamp = Timestamp(int(datetime.now().timestamp() * 1000)) self.processor_status = Status() self._refined_data = {} self.post_hook: Optional[Callable] = None diff --git a/ingestion/src/metadata/data_quality/interface/pandas/pandas_test_suite_interface.py b/ingestion/src/metadata/data_quality/interface/pandas/pandas_test_suite_interface.py index d426f2c87d39..6f3573857354 100644 --- a/ingestion/src/metadata/data_quality/interface/pandas/pandas_test_suite_interface.py +++ b/ingestion/src/metadata/data_quality/interface/pandas/pandas_test_suite_interface.py @@ -13,7 +13,7 @@ Interfaces with database for all database engine supporting sqlalchemy abstraction layer """ -from datetime import datetime, timezone +from datetime import datetime from typing import Optional from metadata.data_quality.interface.test_suite_interface import TestSuiteInterface @@ -94,7 +94,7 @@ def run_test_case( test_handler = TestHandler( self.dfs, test_case=test_case, - execution_date=int(datetime.now(tz=timezone.utc).timestamp() * 1000), + execution_date=int(datetime.now().timestamp() * 1000), ) return Validator(validator_obj=test_handler).validate() diff --git a/ingestion/src/metadata/data_quality/interface/sqlalchemy/sqa_test_suite_interface.py b/ingestion/src/metadata/data_quality/interface/sqlalchemy/sqa_test_suite_interface.py index a0cec41d7734..5e358b9eeea4 100644 --- a/ingestion/src/metadata/data_quality/interface/sqlalchemy/sqa_test_suite_interface.py +++ b/ingestion/src/metadata/data_quality/interface/sqlalchemy/sqa_test_suite_interface.py @@ -14,7 +14,7 @@ supporting sqlalchemy abstraction layer """ -from datetime import datetime, timezone +from datetime import datetime from typing import Optional, Union from sqlalchemy.orm import DeclarativeMeta @@ -169,7 +169,7 @@ def run_test_case( test_handler = TestHandler( self.runner, test_case=test_case, - execution_date=int(datetime.now(tz=timezone.utc).timestamp() * 1000), + execution_date=int(datetime.now().timestamp() * 1000), ) return Validator(validator_obj=test_handler).validate() diff --git a/ingestion/src/metadata/great_expectations/action.py b/ingestion/src/metadata/great_expectations/action.py index e95a86827c35..f7c4a8714d23 100644 --- a/ingestion/src/metadata/great_expectations/action.py +++ b/ingestion/src/metadata/great_expectations/action.py @@ -17,7 +17,7 @@ """ import logging import traceback -from datetime import datetime, timezone +from datetime import datetime from typing import Dict, List, Optional, Union, cast from great_expectations.checkpoint.actions import ValidationAction @@ -423,9 +423,7 @@ def _handle_test_case( self.ometa_conn.add_test_case_results( test_results=TestCaseResult( - timestamp=Timestamp( - int(datetime.now(tz=timezone.utc).timestamp() * 1000) - ), + timestamp=Timestamp(int(datetime.now().timestamp() * 1000)), testCaseStatus=TestCaseStatus.Success if result["success"] else TestCaseStatus.Failed, diff --git a/ingestion/src/metadata/ingestion/bulksink/metadata_usage.py b/ingestion/src/metadata/ingestion/bulksink/metadata_usage.py index f7706f352025..3b8c1551fa93 100644 --- a/ingestion/src/metadata/ingestion/bulksink/metadata_usage.py +++ b/ingestion/src/metadata/ingestion/bulksink/metadata_usage.py @@ -21,7 +21,7 @@ import os import shutil import traceback -from datetime import datetime, timezone +from datetime import datetime from pathlib import Path from typing import List, Optional @@ -39,6 +39,7 @@ StackTraceError, ) from metadata.generated.schema.entity.teams.user import User +from metadata.generated.schema.type.basic import Timestamp from metadata.generated.schema.type.lifeCycle import AccessDetails, LifeCycle from metadata.generated.schema.type.tableUsageCount import TableColumn, TableUsageCount from metadata.generated.schema.type.usageRequest import UsageRequest @@ -282,9 +283,7 @@ def __get_table_joins( Method to get Table Joins """ # TODO: Clean up how we are passing dates from query parsing to here to use timestamps instead of strings - start_date = datetime.fromtimestamp(int(table_usage.date) / 1000).replace( - tzinfo=timezone.utc - ) + start_date = datetime.fromtimestamp(int(table_usage.date) / 1000) table_joins: TableJoins = TableJoins( columnJoins=[], directTableJoins=[], startDate=start_date ) @@ -374,7 +373,7 @@ def _get_table_life_cycle_data( query_type = get_query_type(create_query=create_query) if query_type: access_details = AccessDetails( - timestamp=create_query.queryDate.root, + timestamp=Timestamp(create_query.queryDate.root), accessedBy=user, accessedByAProcess=process_user, ) diff --git a/ingestion/src/metadata/ingestion/connections/test_connections.py b/ingestion/src/metadata/ingestion/connections/test_connections.py index 557be02fd0c5..96794d0070eb 100644 --- a/ingestion/src/metadata/ingestion/connections/test_connections.py +++ b/ingestion/src/metadata/ingestion/connections/test_connections.py @@ -13,7 +13,7 @@ creating a service """ import traceback -from datetime import datetime, timezone +from datetime import datetime from functools import partial from typing import Callable, List, Optional @@ -148,14 +148,14 @@ def _test_connection_steps_automation_workflow( break test_connection_result.lastUpdatedAt = Timestamp( - int(datetime.now(timezone.utc).timestamp() * 1000) + int(datetime.now().timestamp() * 1000) ) metadata.patch_automation_workflow_response( automation_workflow, test_connection_result, WorkflowStatus.Running ) test_connection_result.lastUpdatedAt = Timestamp( - int(datetime.now(timezone.utc).timestamp() * 1000) + int(datetime.now().timestamp() * 1000) ) test_connection_result.status = ( @@ -174,7 +174,7 @@ def _test_connection_steps_automation_workflow( f"Wild error happened while testing the connection in the workflow - {err}" ) logger.debug(traceback.format_exc()) - test_connection_result.lastUpdatedAt = datetime.now(tz=timezone.utc).timestamp() + test_connection_result.lastUpdatedAt = datetime.now().timestamp() metadata.create_or_update( CreateWorkflowRequest( name=automation_workflow.name, diff --git a/ingestion/src/metadata/ingestion/ometa/auth_provider.py b/ingestion/src/metadata/ingestion/ometa/auth_provider.py index 5ae214f87b05..03fb80364bd7 100644 --- a/ingestion/src/metadata/ingestion/ometa/auth_provider.py +++ b/ingestion/src/metadata/ingestion/ometa/auth_provider.py @@ -14,7 +14,7 @@ import os.path from abc import ABCMeta, abstractmethod from dataclasses import dataclass -from datetime import datetime, timezone +from datetime import datetime from dateutil.relativedelta import relativedelta @@ -85,7 +85,7 @@ def __init__(self, config: OpenMetadataConnection): self.config = config self.security_config: OpenMetadataJWTClientConfig = self.config.securityConfig self.jwt_token = None - self.expiry = datetime.now(tz=timezone.utc) - relativedelta(years=1) + self.expiry = datetime.now() - relativedelta(years=1) @classmethod def create(cls, config: OpenMetadataConnection): diff --git a/ingestion/src/metadata/ingestion/ometa/mixins/tests_mixin.py b/ingestion/src/metadata/ingestion/ometa/mixins/tests_mixin.py index 215ce7d8ef78..142f529bfa91 100644 --- a/ingestion/src/metadata/ingestion/ometa/mixins/tests_mixin.py +++ b/ingestion/src/metadata/ingestion/ometa/mixins/tests_mixin.py @@ -15,7 +15,7 @@ """ import traceback -from datetime import datetime, timezone +from datetime import datetime from typing import List, Optional, Type, Union from urllib.parse import quote from uuid import UUID @@ -87,7 +87,7 @@ def get_or_create_test_suite( test_suite_name: str, test_suite_description: Optional[ str - ] = f"Test Suite created on {datetime.now(timezone.utc).strftime('%Y-%m-%d')}", + ] = f"Test Suite created on {datetime.now().strftime('%Y-%m-%d')}", ) -> TestSuite: """Get or create a TestSuite diff --git a/ingestion/src/metadata/ingestion/source/dashboard/looker/metadata.py b/ingestion/src/metadata/ingestion/source/dashboard/looker/metadata.py index a5235cf8caf1..e622d87b8e1d 100644 --- a/ingestion/src/metadata/ingestion/source/dashboard/looker/metadata.py +++ b/ingestion/src/metadata/ingestion/source/dashboard/looker/metadata.py @@ -22,7 +22,7 @@ import copy import os import traceback -from datetime import datetime, timezone +from datetime import datetime from pathlib import Path from typing import Dict, Iterable, List, Optional, Sequence, Set, Type, Union, cast @@ -168,7 +168,7 @@ def __init__( metadata: OpenMetadata, ): super().__init__(config, metadata) - self.today = datetime.now(tz=timezone.utc).strftime("%Y-%m-%d") + self.today = datetime.now().strftime("%Y-%m-%d") self._explores_cache = {} self._repo_credentials: Optional[ReadersCredentials] = None diff --git a/ingestion/src/metadata/ingestion/source/dashboard/metabase/client.py b/ingestion/src/metadata/ingestion/source/dashboard/metabase/client.py index 9204adc58463..2e6fe4c03b72 100644 --- a/ingestion/src/metadata/ingestion/source/dashboard/metabase/client.py +++ b/ingestion/src/metadata/ingestion/source/dashboard/metabase/client.py @@ -55,7 +55,7 @@ def _get_metabase_session(self) -> str: if self.config.password: params[PASSWORD_HEADER] = self.config.password.get_secret_value() self.resp = requests.post( - f"{self.config.hostPort}/{API_VERSION}/session/", + f"{self.config.hostPort}{API_VERSION}/session/", data=json.dumps(params), headers=SESSION_HEADERS, timeout=DEFAULT_TIMEOUT, diff --git a/ingestion/src/metadata/ingestion/source/dashboard/tableau/connection.py b/ingestion/src/metadata/ingestion/source/dashboard/tableau/connection.py index 7b0670bc6e51..8feef66c0e0e 100644 --- a/ingestion/src/metadata/ingestion/source/dashboard/tableau/connection.py +++ b/ingestion/src/metadata/ingestion/source/dashboard/tableau/connection.py @@ -119,8 +119,8 @@ def build_server_config(connection: TableauConnection) -> Dict[str, Dict[str, An """ tableau_server_config = { f"{connection.env}": { - "server": connection.hostPort, - "api_version": connection.apiVersion, + "server": str(connection.hostPort), + "api_version": str(connection.apiVersion), "site_name": connection.siteName if connection.siteName else "", "site_url": connection.siteUrl if connection.siteUrl else "", } diff --git a/ingestion/src/metadata/ingestion/source/database/bigquery/connection.py b/ingestion/src/metadata/ingestion/source/database/bigquery/connection.py index 095c60dd8be5..a6a19aac760b 100644 --- a/ingestion/src/metadata/ingestion/source/database/bigquery/connection.py +++ b/ingestion/src/metadata/ingestion/source/database/bigquery/connection.py @@ -13,7 +13,7 @@ Source connection handler """ import os -from datetime import datetime, timezone +from datetime import datetime from functools import partial from typing import Optional @@ -145,7 +145,7 @@ def test_connection_inner(engine): engine=engine, statement=BIGQUERY_TEST_STATEMENT.format( region=service_connection.usageLocation, - creation_date=datetime.now(tz=timezone.utc).strftime("%Y-%m-%d"), + creation_date=datetime.now().strftime("%Y-%m-%d"), ), ), } diff --git a/ingestion/src/metadata/ingestion/source/database/bigquery/metadata.py b/ingestion/src/metadata/ingestion/source/database/bigquery/metadata.py index 3a69911bf7c7..b3e3624f7998 100644 --- a/ingestion/src/metadata/ingestion/source/database/bigquery/metadata.py +++ b/ingestion/src/metadata/ingestion/source/database/bigquery/metadata.py @@ -369,7 +369,7 @@ def yield_tag( tags=[value], classification_name=key, tag_description="Bigquery Dataset Label", - classification_description="", + classification_description="BigQuery Dataset Classification", include_tags=self.source_config.includeTags, ) # Fetching policy tags on the column level @@ -389,7 +389,7 @@ def yield_tag( tags=[tag.display_name for tag in policy_tags], classification_name=taxonomy.display_name, tag_description="Bigquery Policy Tag", - classification_description="", + classification_description="BigQuery Policy Classification", include_tags=self.source_config.includeTags, ) except Exception as exc: @@ -528,7 +528,7 @@ def yield_table_tags(self, table_name_and_type: Tuple[str, str]): tags=[value], classification_name=key, tag_description="Bigquery Table Label", - classification_description="", + classification_description="BigQuery Table Classification", include_tags=self.source_config.includeTags, ) diff --git a/ingestion/src/metadata/ingestion/source/database/databricks/lineage.py b/ingestion/src/metadata/ingestion/source/database/databricks/lineage.py index a00a8980c882..543351eb45f4 100644 --- a/ingestion/src/metadata/ingestion/source/database/databricks/lineage.py +++ b/ingestion/src/metadata/ingestion/source/database/databricks/lineage.py @@ -12,7 +12,7 @@ Databricks lineage module """ import traceback -from datetime import datetime, timezone +from datetime import datetime from typing import Iterator from metadata.generated.schema.type.basic import DateTime @@ -44,7 +44,7 @@ def yield_table_query(self) -> Iterator[TableQuery]: userName=row.get("user_name"), startTime=row.get("query_start_time_ms"), endTime=row.get("execution_end_time_ms"), - analysisDate=DateTime(datetime.now(tz=timezone.utc)), + analysisDate=DateTime(datetime.now()), serviceName=self.config.serviceName, ) except Exception as exc: diff --git a/ingestion/src/metadata/ingestion/source/database/databricks/metadata.py b/ingestion/src/metadata/ingestion/source/database/databricks/metadata.py index c04ff31993be..192084db6743 100644 --- a/ingestion/src/metadata/ingestion/source/database/databricks/metadata.py +++ b/ingestion/src/metadata/ingestion/source/database/databricks/metadata.py @@ -68,8 +68,8 @@ logger = ingestion_logger() -DATABRICKS_TAG = "DATABRICK TAG" -DATABRICKS_TAG_CLASSIFICATION = "DATABRICK TAG CLASSIFICATION" +DATABRICKS_TAG = "DATABRICKS TAG" +DATABRICKS_TAG_CLASSIFICATION = "DATABRICKS TAG CLASSIFICATION" DEFAULT_TAG_VALUE = "NONE" diff --git a/ingestion/src/metadata/ingestion/source/database/databricks/usage.py b/ingestion/src/metadata/ingestion/source/database/databricks/usage.py index 6c74976fff0a..b2b31f259b03 100644 --- a/ingestion/src/metadata/ingestion/source/database/databricks/usage.py +++ b/ingestion/src/metadata/ingestion/source/database/databricks/usage.py @@ -12,7 +12,7 @@ Databricks usage module """ import traceback -from datetime import datetime, timezone +from datetime import datetime from typing import Iterable from metadata.generated.schema.type.basic import DateTime @@ -49,7 +49,7 @@ def yield_table_queries(self) -> Iterable[TableQuery]: userName=row.get("user_name"), startTime=row.get("query_start_time_ms"), endTime=row.get("execution_end_time_ms"), - analysisDate=DateTime(datetime.now(tz=timezone.utc)), + analysisDate=DateTime(datetime.now()), serviceName=self.config.serviceName, duration=row.get("duration") if row.get("duration") diff --git a/ingestion/src/metadata/ingestion/source/database/incremental_metadata_extraction.py b/ingestion/src/metadata/ingestion/source/database/incremental_metadata_extraction.py index 63071644cfa1..22e5d2e59a8e 100644 --- a/ingestion/src/metadata/ingestion/source/database/incremental_metadata_extraction.py +++ b/ingestion/src/metadata/ingestion/source/database/incremental_metadata_extraction.py @@ -70,7 +70,7 @@ def __init__( def _calculate_pipeline_status_parameters(self) -> Tuple[int, int]: """Calculate the needed 'start' and 'end' parameters based on the 'lookbackDays'.""" - now = datetime.now(tz=timezone.utc) + now = datetime.now() # We multiply the value by 1000 because our backend uses epoch_milliseconds instead of epoch_seconds. start = int( diff --git a/ingestion/src/metadata/ingestion/source/database/life_cycle_query_mixin.py b/ingestion/src/metadata/ingestion/source/database/life_cycle_query_mixin.py index b4a8ca81533b..546911666f57 100644 --- a/ingestion/src/metadata/ingestion/source/database/life_cycle_query_mixin.py +++ b/ingestion/src/metadata/ingestion/source/database/life_cycle_query_mixin.py @@ -27,6 +27,7 @@ from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline import ( DatabaseServiceMetadataPipeline, ) +from metadata.generated.schema.type.basic import Timestamp from metadata.generated.schema.type.lifeCycle import AccessDetails, LifeCycle from metadata.ingestion.api.models import Either, Entity from metadata.ingestion.api.status import Status @@ -103,8 +104,12 @@ def get_life_cycle_data( if life_cycle_data: life_cycle = LifeCycle( created=AccessDetails( - timestamp=convert_timestamp_to_milliseconds( - life_cycle_data.created_at.timestamp() + timestamp=Timestamp( + int( + convert_timestamp_to_milliseconds( + life_cycle_data.created_at.timestamp() + ) + ) ) ) ) diff --git a/ingestion/src/metadata/ingestion/source/database/postgres/lineage.py b/ingestion/src/metadata/ingestion/source/database/postgres/lineage.py index 9b72272088ef..754f222c689b 100644 --- a/ingestion/src/metadata/ingestion/source/database/postgres/lineage.py +++ b/ingestion/src/metadata/ingestion/source/database/postgres/lineage.py @@ -12,7 +12,7 @@ Postgres lineage module """ import traceback -from datetime import datetime, timezone +from datetime import datetime from typing import Iterable from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest @@ -84,7 +84,7 @@ def process_table_query(self) -> Iterable[TableQuery]: yield TableQuery( query=row["query_text"], userName=row["usename"], - analysisDate=DateTime(datetime.now(tz=timezone.utc)), + analysisDate=DateTime(datetime.now()), aborted=self.get_aborted_status(row), databaseName=self.get_database_name(row), serviceName=self.config.serviceName, diff --git a/ingestion/src/metadata/ingestion/source/database/postgres/usage.py b/ingestion/src/metadata/ingestion/source/database/postgres/usage.py index 4014d2cede32..d6603b5a4896 100644 --- a/ingestion/src/metadata/ingestion/source/database/postgres/usage.py +++ b/ingestion/src/metadata/ingestion/source/database/postgres/usage.py @@ -12,7 +12,7 @@ Postgres usage module """ import traceback -from datetime import datetime, timezone +from datetime import datetime from typing import Iterable from metadata.generated.schema.type.basic import DateTime @@ -53,7 +53,7 @@ def process_table_query(self) -> Iterable[TableQueries]: TableQuery( query=row["query_text"], userName=row["usename"], - analysisDate=DateTime(datetime.now(tz=timezone.utc)), + analysisDate=DateTime(datetime.now()), aborted=self.get_aborted_status(row), databaseName=self.get_database_name(row), serviceName=self.config.serviceName, diff --git a/ingestion/src/metadata/ingestion/source/database/sample_data.py b/ingestion/src/metadata/ingestion/source/database/sample_data.py index 96161c00adb6..3951410832f2 100644 --- a/ingestion/src/metadata/ingestion/source/database/sample_data.py +++ b/ingestion/src/metadata/ingestion/source/database/sample_data.py @@ -18,7 +18,7 @@ import time import traceback from collections import namedtuple -from datetime import datetime, timedelta, timezone +from datetime import datetime, timedelta from typing import Any, Dict, Iterable, List, Optional, Union from pydantic import ValidationError @@ -1343,10 +1343,7 @@ def ingest_profiles(self) -> Iterable[Either[OMetaTableProfileSampleData]]: customMetrics=profile.get("customMetrics"), timestamp=Timestamp( int( - ( - datetime.now(tz=timezone.utc) - - timedelta(days=days) - ).timestamp() + (datetime.now() - timedelta(days=days)).timestamp() * 1000 ) ), @@ -1356,8 +1353,7 @@ def ingest_profiles(self) -> Iterable[Either[OMetaTableProfileSampleData]]: timestamp=Timestamp( int( ( - datetime.now(tz=timezone.utc) - - timedelta(days=days) + datetime.now() - timedelta(days=days) ).timestamp() * 1000 ) @@ -1371,7 +1367,7 @@ def ingest_profiles(self) -> Iterable[Either[OMetaTableProfileSampleData]]: timestamp=Timestamp( int( ( - datetime.now(tz=timezone.utc) + datetime.now() - timedelta( days=days, hours=random.randint(0, 24) ) @@ -1519,10 +1515,7 @@ def ingest_test_case_results(self) -> Iterable[Either[OMetaTestCaseResultsSample test_case_results=TestCaseResult( timestamp=Timestamp( int( - ( - datetime.now(tz=timezone.utc) - - timedelta(days=days) - ).timestamp() + (datetime.now() - timedelta(days=days)).timestamp() * 1000 ) ), @@ -1570,10 +1563,7 @@ def ingest_data_insights(self) -> Iterable[Either[OMetaDataInsightSample]]: reportDataType=report_datum["reportDataType"], timestamp=Timestamp( root=int( - ( - datetime.now(tz=timezone.utc) - timedelta(days=i) - ).timestamp() - * 1000 + (datetime.now() - timedelta(days=i)).timestamp() * 1000 ) ), data=report_datum["data"], @@ -1592,7 +1582,7 @@ def ingest_life_cycle(self) -> Iterable[Either[OMetaLifeCycleData]]: int( convert_timestamp_to_milliseconds( ( - datetime.now(tz=timezone.utc) + datetime.now() - timedelta(days=life_cycle["created"]["days"]) ).timestamp() ) @@ -1606,7 +1596,7 @@ def ingest_life_cycle(self) -> Iterable[Either[OMetaLifeCycleData]]: int( convert_timestamp_to_milliseconds( ( - datetime.now(tz=timezone.utc) + datetime.now() - timedelta(days=life_cycle["updated"]["days"]) ).timestamp() ) @@ -1620,7 +1610,7 @@ def ingest_life_cycle(self) -> Iterable[Either[OMetaLifeCycleData]]: int( convert_timestamp_to_milliseconds( ( - datetime.now(tz=timezone.utc) + datetime.now() - timedelta(days=life_cycle["accessed"]["days"]) ).timestamp() ) diff --git a/ingestion/src/metadata/ingestion/source/database/sample_usage.py b/ingestion/src/metadata/ingestion/source/database/sample_usage.py index 89c49bafd8e2..41b30ca181f8 100644 --- a/ingestion/src/metadata/ingestion/source/database/sample_usage.py +++ b/ingestion/src/metadata/ingestion/source/database/sample_usage.py @@ -13,7 +13,7 @@ """ import csv import json -from datetime import datetime, timezone +from datetime import datetime from typing import Dict, Iterable, Optional from metadata.generated.schema.entity.services.connections.database.customDatabaseConnection import ( @@ -47,7 +47,7 @@ class SampleUsageSource(UsageSource): def __init__(self, config: WorkflowSource, metadata: OpenMetadata): super().__init__(config, metadata, False) - self.analysis_date = DateTime(datetime.now(tz=timezone.utc)) + self.analysis_date = DateTime(datetime.now()) sample_data_folder = self.service_connection.connectionOptions.root.get( "sampleDataFolder" diff --git a/ingestion/src/metadata/ingestion/source/database/sas/metadata.py b/ingestion/src/metadata/ingestion/source/database/sas/metadata.py index 498e56368739..082280d9736c 100644 --- a/ingestion/src/metadata/ingestion/source/database/sas/metadata.py +++ b/ingestion/src/metadata/ingestion/source/database/sas/metadata.py @@ -122,7 +122,7 @@ def __init__(self, config: WorkflowSource, metadata: OpenMetadata): self.databases = None self.database_schemas = None - self.timestamp = Timestamp(int(datetime.now(timezone.utc).timestamp() * 1000)) + self.timestamp = Timestamp(int(datetime.now().timestamp() * 1000)) @classmethod def create( @@ -572,11 +572,17 @@ def create_table_entity(self, table) -> Iterable[Either[CreateTableRequest]]: ): return + raw_create_date: Optional[datetime] = table_entity_instance.get( + "creationTimeStamp" + ) + if raw_create_date: + raw_create_date = raw_create_date.replace(tzinfo=timezone.utc) + # create Profiles & Data Quality Column table_profile_request = CreateTableProfileRequest( tableProfile=TableProfile( timestamp=self.timestamp, - createDateTime=table_entity_instance["creationTimeStamp"], + createDateTime=raw_create_date, rowCount=int(table_extension.get("rowCount", 0)), columnCount=int(table_extension.get("columnCount", 0)), sizeInByte=int(table_extension.get("dataSize", 0)), diff --git a/ingestion/src/metadata/ingestion/source/database/stored_procedures_mixin.py b/ingestion/src/metadata/ingestion/source/database/stored_procedures_mixin.py index e5ad89fc21d1..0c58e50ad2c4 100644 --- a/ingestion/src/metadata/ingestion/source/database/stored_procedures_mixin.py +++ b/ingestion/src/metadata/ingestion/source/database/stored_procedures_mixin.py @@ -18,7 +18,7 @@ from datetime import datetime from typing import Dict, Iterable, List, Optional, Union -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field from sqlalchemy.engine import Engine from metadata.generated.schema.api.data.createQuery import CreateQueryRequest @@ -53,18 +53,17 @@ class QueryByProcedure(BaseModel): procedure_name: str = Field(None, alias="PROCEDURE_NAME") query_type: str = Field(..., alias="QUERY_TYPE") - query_database_name: str = Field(None, alias="QUERY_DATABASE_NAME") - query_schema_name: str = Field(None, alias="QUERY_SCHEMA_NAME") + query_database_name: Optional[str] = Field(None, alias="QUERY_DATABASE_NAME") + query_schema_name: Optional[str] = Field(None, alias="QUERY_SCHEMA_NAME") procedure_text: str = Field(..., alias="PROCEDURE_TEXT") procedure_start_time: datetime = Field(..., alias="PROCEDURE_START_TIME") procedure_end_time: datetime = Field(..., alias="PROCEDURE_END_TIME") - query_start_time: Optional[datetime] = Field(..., alias="QUERY_START_TIME") + query_start_time: Optional[datetime] = Field(None, alias="QUERY_START_TIME") query_duration: Optional[float] = Field(None, alias="QUERY_DURATION") query_text: str = Field(..., alias="QUERY_TEXT") query_user_name: Optional[str] = Field(None, alias="QUERY_USER_NAME") - class Config: - populate_by_name = True + model_config = ConfigDict(populate_by_name=True) class StoredProcedureMixin(ABC): diff --git a/ingestion/src/metadata/ingestion/source/database/usage_source.py b/ingestion/src/metadata/ingestion/source/database/usage_source.py index eb4ed6c08a69..769f1c125edc 100644 --- a/ingestion/src/metadata/ingestion/source/database/usage_source.py +++ b/ingestion/src/metadata/ingestion/source/database/usage_source.py @@ -14,7 +14,7 @@ import csv import traceback from abc import ABC -from datetime import datetime, timedelta, timezone +from datetime import datetime, timedelta from typing import Iterable from metadata.generated.schema.type.basic import DateTime @@ -45,18 +45,13 @@ def yield_table_queries_from_logs(self) -> Iterable[TableQuery]: for record in csv.DictReader(fin): query_dict = dict(record) - # Prepare the timezone aware analysis date - if query_dict.get("start_time"): - # We allow reading the start_time without timezone for simplicity for users - timestamp = datetime.strptime( + analysis_date = ( + datetime.utcnow() + if not query_dict.get("start_time") + else datetime.strptime( query_dict.get("start_time"), "%Y-%m-%d %H:%M:%S.%f" - ).timestamp() - analysis_date = datetime.fromtimestamp( - timestamp, tz=timezone.utc ) - else: - analysis_date = datetime.now(tz=timezone.utc) - + ) query_list.append( TableQuery( query=query_dict["query_text"], diff --git a/ingestion/src/metadata/ingestion/source/pipeline/kafkaconnect/metadata.py b/ingestion/src/metadata/ingestion/source/pipeline/kafkaconnect/metadata.py index df3f8cb46304..086c3030accc 100644 --- a/ingestion/src/metadata/ingestion/source/pipeline/kafkaconnect/metadata.py +++ b/ingestion/src/metadata/ingestion/source/pipeline/kafkaconnect/metadata.py @@ -12,7 +12,7 @@ KafkaConnect source to extract metadata from OM UI """ import traceback -from datetime import datetime, timezone +from datetime import datetime from typing import Iterable, Optional from metadata.generated.schema.api.data.createPipeline import CreatePipelineRequest @@ -291,7 +291,7 @@ def yield_pipeline_status( pipeline_details.status, StatusType.Pending ), taskStatus=task_status, - timestamp=Timestamp(datetime_to_ts(datetime.now(tz=timezone.utc))) + timestamp=Timestamp(datetime_to_ts(datetime.now())) # Kafka connect doesn't provide any details with exec time ) diff --git a/ingestion/src/metadata/ingestion/source/storage/s3/metadata.py b/ingestion/src/metadata/ingestion/source/storage/s3/metadata.py index 3f04c9445786..e82df1002ff7 100644 --- a/ingestion/src/metadata/ingestion/source/storage/s3/metadata.py +++ b/ingestion/src/metadata/ingestion/source/storage/s3/metadata.py @@ -12,7 +12,7 @@ import json import secrets import traceback -from datetime import datetime, timedelta, timezone +from datetime import datetime, timedelta from enum import Enum from typing import Dict, Iterable, List, Optional @@ -301,9 +301,9 @@ def _fetch_metric(self, bucket_name: str, metric: S3Metric) -> float: }, }, ], - StartTime=datetime.now(tz=timezone.utc) - timedelta(days=2), + StartTime=datetime.now() - timedelta(days=2), # metrics generated daily, ensure there is at least 1 entry - EndTime=datetime.now(tz=timezone.utc), + EndTime=datetime.now(), ScanBy="TimestampDescending", ) if raw_result["MetricDataResults"]: diff --git a/ingestion/src/metadata/profiler/interface/nosql/profiler_interface.py b/ingestion/src/metadata/profiler/interface/nosql/profiler_interface.py index c23fdd267672..4a206bf05f92 100644 --- a/ingestion/src/metadata/profiler/interface/nosql/profiler_interface.py +++ b/ingestion/src/metadata/profiler/interface/nosql/profiler_interface.py @@ -16,7 +16,7 @@ """ import traceback from collections import defaultdict -from datetime import datetime, timezone +from datetime import datetime from typing import Dict, List, Optional, Type from sqlalchemy import Column @@ -212,9 +212,7 @@ def get_all_metrics( profile_results["columns"][column].update( { "name": column, - "timestamp": int( - datetime.now(tz=timezone.utc).timestamp() * 1000 - ), + "timestamp": int(datetime.now().timestamp() * 1000), **profile, } ) diff --git a/ingestion/src/metadata/profiler/interface/pandas/profiler_interface.py b/ingestion/src/metadata/profiler/interface/pandas/profiler_interface.py index 89ebf7f6435c..8d6c3beffe2c 100644 --- a/ingestion/src/metadata/profiler/interface/pandas/profiler_interface.py +++ b/ingestion/src/metadata/profiler/interface/pandas/profiler_interface.py @@ -17,7 +17,7 @@ import traceback from collections import defaultdict from copy import deepcopy -from datetime import datetime, timezone +from datetime import datetime from typing import Dict, List, Optional from sqlalchemy import Column @@ -403,9 +403,7 @@ def get_all_metrics( profile_results["columns"][column].update( { "name": column, - "timestamp": int( - datetime.now(tz=timezone.utc).timestamp() * 1000 - ), + "timestamp": int(datetime.now().timestamp() * 1000), **profile, } ) diff --git a/ingestion/src/metadata/profiler/interface/sqlalchemy/profiler_interface.py b/ingestion/src/metadata/profiler/interface/sqlalchemy/profiler_interface.py index da59e2063511..442c938abfa3 100644 --- a/ingestion/src/metadata/profiler/interface/sqlalchemy/profiler_interface.py +++ b/ingestion/src/metadata/profiler/interface/sqlalchemy/profiler_interface.py @@ -19,7 +19,7 @@ import threading import traceback from collections import defaultdict -from datetime import datetime, timezone +from datetime import datetime from typing import Dict, List, Optional from sqlalchemy import Column, inspect, text @@ -506,9 +506,7 @@ def get_all_metrics( profile_results["columns"][column].update( { "name": column, - "timestamp": int( - datetime.now(tz=timezone.utc).timestamp() * 1000 - ), + "timestamp": int(datetime.now().timestamp() * 1000), **profile, } ) diff --git a/ingestion/src/metadata/profiler/processor/core.py b/ingestion/src/metadata/profiler/processor/core.py index f3b666ff1d55..c77c14a04aaa 100644 --- a/ingestion/src/metadata/profiler/processor/core.py +++ b/ingestion/src/metadata/profiler/processor/core.py @@ -106,9 +106,7 @@ def __init__( self.include_columns = include_columns self.exclude_columns = exclude_columns self._metrics = metrics - self._profile_ts = Timestamp( - int(datetime.now(tz=timezone.utc).timestamp() * 1000) - ) + self._profile_ts = Timestamp(int(datetime.now().timestamp() * 1000)) self.profile_sample_config = self.profiler_interface.profile_sample_config self.metric_filter = MetricFilter( diff --git a/ingestion/src/metadata/profiler/processor/sample_data_handler.py b/ingestion/src/metadata/profiler/processor/sample_data_handler.py index 674633d16f77..733c66c8e8a7 100644 --- a/ingestion/src/metadata/profiler/processor/sample_data_handler.py +++ b/ingestion/src/metadata/profiler/processor/sample_data_handler.py @@ -13,7 +13,7 @@ """ import json import traceback -from datetime import datetime, timezone +from datetime import datetime from functools import singledispatch from io import BytesIO @@ -65,7 +65,7 @@ def _get_object_key( ) if not overwrite_data: file_name = file_name.replace( - ".parquet", f"_{datetime.now(tz=timezone.utc).strftime('%Y_%m_%d')}.parquet" + ".parquet", f"_{datetime.now().strftime('%Y_%m_%d')}.parquet" ) if prefix: return f"{clean_uri(prefix)}/{file_name}" diff --git a/ingestion/src/metadata/utils/tag_utils.py b/ingestion/src/metadata/utils/tag_utils.py index 89407194c9b1..5bcde6b0e537 100644 --- a/ingestion/src/metadata/utils/tag_utils.py +++ b/ingestion/src/metadata/utils/tag_utils.py @@ -49,8 +49,8 @@ def get_ometa_tag_and_classification( tags: List[str], classification_name: str, - tag_description: Optional[str] = None, - classification_description: Optional[str] = None, + tag_description: str, + classification_description: str, include_tags: bool = True, tag_fqn: Optional[FullyQualifiedEntityName] = None, ) -> Iterable[Either[OMetaTagAndClassification]]: @@ -64,9 +64,7 @@ def get_ometa_tag_and_classification( fqn=tag_fqn, classification_request=CreateClassificationRequest( name=EntityName(classification_name), - description=Markdown(classification_description) - if classification_description - else None, + description=Markdown(classification_description), ), tag_request=CreateTagRequest( classification=FullyQualifiedEntityName(classification_name), diff --git a/ingestion/src/metadata/workflow/base.py b/ingestion/src/metadata/workflow/base.py index 9bef2dfa2ace..dbde3bbe1c70 100644 --- a/ingestion/src/metadata/workflow/base.py +++ b/ingestion/src/metadata/workflow/base.py @@ -14,7 +14,7 @@ import uuid from abc import ABC, abstractmethod -from datetime import datetime, timezone +from datetime import datetime from typing import Any, Dict, List, Optional, TypeVar, Union from metadata.generated.schema.api.services.ingestionPipelines.createIngestionPipeline import ( @@ -92,7 +92,7 @@ def __init__( self.service_type = service_type self._timer: Optional[RepeatedTimer] = None self._ingestion_pipeline: Optional[IngestionPipeline] = None - self._start_ts = datetime_to_ts(datetime.now(tz=timezone.utc)) + self._start_ts = datetime_to_ts(datetime.now()) self._execution_time_tracker = ExecutionTimeTracker( log_level == LogLevels.DEBUG ) diff --git a/ingestion/src/metadata/workflow/workflow_status_mixin.py b/ingestion/src/metadata/workflow/workflow_status_mixin.py index 213c04e70ef6..9772857a1a8a 100644 --- a/ingestion/src/metadata/workflow/workflow_status_mixin.py +++ b/ingestion/src/metadata/workflow/workflow_status_mixin.py @@ -13,7 +13,7 @@ """ import traceback import uuid -from datetime import datetime, timezone +from datetime import datetime from typing import Optional, Tuple from metadata.config.common import WorkflowExecutionError @@ -100,7 +100,7 @@ def set_ingestion_pipeline_status( else: # if workflow is ended then update the end date in status pipeline_status.endDate = Timestamp( - int(datetime.now(tz=timezone.utc).timestamp() * 1000) + int(datetime.now().timestamp() * 1000) ) pipeline_status.pipelineState = state diff --git a/ingestion/tests/cli_e2e/base/test_cli.py b/ingestion/tests/cli_e2e/base/test_cli.py index df84c265e676..84ccc4ce30c5 100644 --- a/ingestion/tests/cli_e2e/base/test_cli.py +++ b/ingestion/tests/cli_e2e/base/test_cli.py @@ -93,7 +93,12 @@ def extract_source_status(output) -> Status: output_clean = output_clean_ansi.sub(" ", output_clean) regex = r"[\w] Status:%(log)s(.*?)%(log)s.* Status: .*" % REGEX_AUX output_clean_regex = re.findall(regex, output_clean.strip()) - return Status.parse_obj(literal_eval(output_clean_regex[0].strip())) + try: + return Status.model_validate(literal_eval(output_clean_regex[0].strip())) + except Exception as exc: + raise RuntimeError( + f"Error extracting source status: {exc}. Check the output {output}" + ) @staticmethod def extract_sink_status(output) -> Status: @@ -105,7 +110,12 @@ def extract_sink_status(output) -> Status: r".*OpenMetadata Status:%(log)s(.*?)%(log)sExecution.*Summary.*" % REGEX_AUX ) output_clean_regex = re.findall(regex, output_clean.strip())[0].strip() - return Status.parse_obj(literal_eval(output_clean_regex)) + try: + return Status.model_validate(literal_eval(output_clean_regex)) + except Exception as exc: + raise RuntimeError( + f"Error extracting sink status: {exc}. Check the output {output}" + ) @staticmethod def build_yaml(config_yaml: dict, test_type: E2EType, extra_args: dict): diff --git a/ingestion/tests/cli_e2e/dashboard/tableau/tableau.yaml b/ingestion/tests/cli_e2e/dashboard/tableau/tableau.yaml index 7d1c314921d1..624267c3eb98 100644 --- a/ingestion/tests/cli_e2e/dashboard/tableau/tableau.yaml +++ b/ingestion/tests/cli_e2e/dashboard/tableau/tableau.yaml @@ -11,7 +11,7 @@ source: hostPort: $E2E_TABLEAU_HOST_PORT siteName: $E2E_TABLEAU_SITE siteUrl: $E2E_TABLEAU_SITE - apiVersion: 3.22 + apiVersion: "3.22" sourceConfig: config: type: DashboardMetadata diff --git a/ingestion/tests/cli_e2e/database/bigquery/bigquery.yaml b/ingestion/tests/cli_e2e/database/bigquery/bigquery.yaml index 35f425a8ef2d..e52bf9fdf3fd 100644 --- a/ingestion/tests/cli_e2e/database/bigquery/bigquery.yaml +++ b/ingestion/tests/cli_e2e/database/bigquery/bigquery.yaml @@ -15,7 +15,7 @@ source: privateKeyId: $E2E_BQ_PRIVATE_KEY_ID privateKey: $E2E_BQ_PRIVATE_KEY clientEmail: $E2E_BQ_CLIENT_EMAIL - clientId: $E2E_BQ_CLIENT_ID + clientId: "$E2E_BQ_CLIENT_ID" sourceConfig: config: type: DatabaseMetadata diff --git a/ingestion/tests/cli_e2e/database/bigquery_multiple_project/bigquery_multiple_project.yaml b/ingestion/tests/cli_e2e/database/bigquery_multiple_project/bigquery_multiple_project.yaml index ab961ce816f6..45e61f48e20e 100644 --- a/ingestion/tests/cli_e2e/database/bigquery_multiple_project/bigquery_multiple_project.yaml +++ b/ingestion/tests/cli_e2e/database/bigquery_multiple_project/bigquery_multiple_project.yaml @@ -17,7 +17,7 @@ source: privateKeyId: $E2E_BQ_PRIVATE_KEY_ID privateKey: $E2E_BQ_PRIVATE_KEY clientEmail: $E2E_BQ_CLIENT_EMAIL - clientId: $E2E_BQ_CLIENT_ID + clientId: "$E2E_BQ_CLIENT_ID" sourceConfig: config: type: DatabaseMetadata diff --git a/ingestion/tests/integration/data_insight/producer/test_producers.py b/ingestion/tests/integration/data_insight/producer/test_producers.py index b09e085a74e9..05b3c25ef035 100644 --- a/ingestion/tests/integration/data_insight/producer/test_producers.py +++ b/ingestion/tests/integration/data_insight/producer/test_producers.py @@ -68,7 +68,7 @@ class TestEntityProducer(TestCase): def setUpClass(cls): """test init""" cls.metadata = OpenMetadata( - OpenMetadataConnection.parse_obj( + OpenMetadataConnection.model_validate( data_insight_config["workflowConfig"]["openMetadataServerConfig"] ) ) @@ -106,7 +106,7 @@ def setUpClass(cls): """test init""" cls.event_ids = [] cls.metadata = OpenMetadata( - OpenMetadataConnection.parse_obj( + OpenMetadataConnection.model_validate( data_insight_config["workflowConfig"]["openMetadataServerConfig"] ) ) @@ -142,7 +142,7 @@ def setUpClass(cls): ), ) event = cls.metadata.add_web_analytic_events(create_event) - event = WebAnalyticEventData.parse_obj(event) + event = WebAnalyticEventData.model_validate(event) cls.event_ids.append(event.eventId.root) cls.producer = WebAnalyticsProducer(cls.metadata) diff --git a/ingestion/tests/integration/data_insight/test_data_insight_workflow.py b/ingestion/tests/integration/data_insight/test_data_insight_workflow.py index c963be019b58..9940d1967c2d 100644 --- a/ingestion/tests/integration/data_insight/test_data_insight_workflow.py +++ b/ingestion/tests/integration/data_insight/test_data_insight_workflow.py @@ -128,7 +128,7 @@ class DataInsightWorkflowTests(unittest.TestCase): @classmethod def setUpClass(cls) -> None: cls.metadata = OpenMetadata( - OpenMetadataConnection.parse_obj( + OpenMetadataConnection.model_validate( data_insight_config["workflowConfig"]["openMetadataServerConfig"] ) ) diff --git a/ingestion/tests/integration/data_insight/test_web_analytic_events.py b/ingestion/tests/integration/data_insight/test_web_analytic_events.py index 63bd887d295b..bef6d1769660 100644 --- a/ingestion/tests/integration/data_insight/test_web_analytic_events.py +++ b/ingestion/tests/integration/data_insight/test_web_analytic_events.py @@ -68,7 +68,7 @@ def setUpClass(cls) -> None: """Set up om client for the test class""" cls.metadata = OpenMetadata( - OpenMetadataConnection.parse_obj( + OpenMetadataConnection.model_validate( data_insight_config["workflowConfig"]["openMetadataServerConfig"] ) ) diff --git a/ingestion/tests/integration/ometa/test_ometa_lineage_api.py b/ingestion/tests/integration/ometa/test_ometa_lineage_api.py index 953d8abc81dd..f72f133e1eba 100644 --- a/ingestion/tests/integration/ometa/test_ometa_lineage_api.py +++ b/ingestion/tests/integration/ometa/test_ometa_lineage_api.py @@ -315,5 +315,5 @@ def test_table_datamodel_lineage(self): entity=DashboardDataModel, fqn=self.dashboard_datamodel_entity.fullyQualifiedName.root, ) - entity_lineage = EntityLineage.parse_obj(datamodel_lineage) + entity_lineage = EntityLineage.model_validate(datamodel_lineage) self.assertEqual(from_id, str(entity_lineage.upstreamEdges[0].fromEntity.root)) diff --git a/ingestion/tests/integration/ometa/test_ometa_table_api.py b/ingestion/tests/integration/ometa/test_ometa_table_api.py index cbc661320eba..bc86de63a2fc 100644 --- a/ingestion/tests/integration/ometa/test_ometa_table_api.py +++ b/ingestion/tests/integration/ometa/test_ometa_table_api.py @@ -14,7 +14,7 @@ """ import uuid from copy import deepcopy -from datetime import datetime, timezone +from datetime import datetime from typing import List from unittest import TestCase from unittest.mock import patch @@ -383,24 +383,18 @@ def test_ingest_table_profile_data(self): mean=1.5, sum=2, stddev=None, - timestamp=Timestamp( - root=int(datetime.now(tz=timezone.utc).timestamp()) - ), + timestamp=Timestamp(root=int(datetime.now().timestamp())), ) ] system_profile = [ SystemProfile( - timestamp=Timestamp( - root=int(datetime.now(tz=timezone.utc).timestamp()) - ), + timestamp=Timestamp(root=int(datetime.now().timestamp())), operation=DmlOperationType.INSERT, rowsAffected=11, ), SystemProfile( - timestamp=Timestamp( - root=int(datetime.now(tz=timezone.utc).timestamp()) + 1 - ), + timestamp=Timestamp(root=int(datetime.now().timestamp()) + 1), operation=DmlOperationType.UPDATE, rowsAffected=110, ), diff --git a/ingestion/tests/integration/orm_profiler/system/test_bigquery_system_metrics.py b/ingestion/tests/integration/orm_profiler/system/test_bigquery_system_metrics.py index 99e08b31f623..06227197698d 100644 --- a/ingestion/tests/integration/orm_profiler/system/test_bigquery_system_metrics.py +++ b/ingestion/tests/integration/orm_profiler/system/test_bigquery_system_metrics.py @@ -134,7 +134,9 @@ def setUpClass(cls) -> None: cls.metadata_config_dict = cls.config["workflowConfig"][ "openMetadataServerConfig" ] - cls.metadata_config = OpenMetadataConnection.parse_obj(cls.metadata_config_dict) + cls.metadata_config = OpenMetadataConnection.model_validate( + cls.metadata_config_dict + ) cls.metadata = OpenMetadata(cls.metadata_config) # run the ingestion workflow diff --git a/ingestion/tests/integration/orm_profiler/system/test_redshift_system_metrics.py b/ingestion/tests/integration/orm_profiler/system/test_redshift_system_metrics.py index 95e743cc37f0..734b588e88f5 100644 --- a/ingestion/tests/integration/orm_profiler/system/test_redshift_system_metrics.py +++ b/ingestion/tests/integration/orm_profiler/system/test_redshift_system_metrics.py @@ -101,7 +101,9 @@ def setUpClass(cls) -> None: cls.metadata_config_dict = cls.config["workflowConfig"][ "openMetadataServerConfig" ] - cls.metadata_config = OpenMetadataConnection.parse_obj(cls.metadata_config_dict) + cls.metadata_config = OpenMetadataConnection.model_validate( + cls.metadata_config_dict + ) cls.metadata = OpenMetadata(cls.metadata_config) # run the ingestion workflow diff --git a/ingestion/tests/integration/orm_profiler/system/test_snowflake_system_metrics.py b/ingestion/tests/integration/orm_profiler/system/test_snowflake_system_metrics.py index 278475156e67..d64a08a065d7 100644 --- a/ingestion/tests/integration/orm_profiler/system/test_snowflake_system_metrics.py +++ b/ingestion/tests/integration/orm_profiler/system/test_snowflake_system_metrics.py @@ -125,7 +125,9 @@ def setUpClass(cls) -> None: cls.metadata_config_dict = cls.config["workflowConfig"][ "openMetadataServerConfig" ] - cls.metadata_config = OpenMetadataConnection.parse_obj(cls.metadata_config_dict) + cls.metadata_config = OpenMetadataConnection.model_validate( + cls.metadata_config_dict + ) cls.metadata = OpenMetadata(cls.metadata_config) # run the ingestion workflow diff --git a/ingestion/tests/integration/test_suite/test_e2e_workflow.py b/ingestion/tests/integration/test_suite/test_e2e_workflow.py index 9cd17781ad0c..9cce94d7b1e0 100644 --- a/ingestion/tests/integration/test_suite/test_e2e_workflow.py +++ b/ingestion/tests/integration/test_suite/test_e2e_workflow.py @@ -116,7 +116,7 @@ class TestE2EWorkflow(unittest.TestCase): """e2e test for the workflow""" metadata = OpenMetadata( - OpenMetadataConnection.parse_obj( + OpenMetadataConnection.model_validate( test_suite_config["workflowConfig"]["openMetadataServerConfig"] ) ) diff --git a/ingestion/tests/integration/test_suite/test_registry_names_match_test_definition.py b/ingestion/tests/integration/test_suite/test_registry_names_match_test_definition.py index 96bfcaabe49e..b208fed2c7ae 100644 --- a/ingestion/tests/integration/test_suite/test_registry_names_match_test_definition.py +++ b/ingestion/tests/integration/test_suite/test_registry_names_match_test_definition.py @@ -54,7 +54,7 @@ class TestRegistryNamesMatchTestDefinition(TestCase): """Test the names in the registry match that of the ones in the Test Definition""" metadata = OpenMetadata( - OpenMetadataConnection.parse_obj( + OpenMetadataConnection.model_validate( test_suite_config["workflowConfig"]["openMetadataServerConfig"] ) ) diff --git a/ingestion/tests/integration/test_suite/test_workflow.py b/ingestion/tests/integration/test_suite/test_workflow.py index 0a780662178d..240493bb104d 100644 --- a/ingestion/tests/integration/test_suite/test_workflow.py +++ b/ingestion/tests/integration/test_suite/test_workflow.py @@ -81,7 +81,7 @@ class TestSuiteWorkflowTests(unittest.TestCase): """Main test suite integration tests definition""" metadata = OpenMetadata( - OpenMetadataConnection.parse_obj( + OpenMetadataConnection.model_validate( get_test_suite_config("", "")["workflowConfig"]["openMetadataServerConfig"] ) ) diff --git a/ingestion/tests/unit/profiler/pandas/test_profiler_interface.py b/ingestion/tests/unit/profiler/pandas/test_profiler_interface.py index 5ba20fe62dbc..e3ae5a9177d4 100644 --- a/ingestion/tests/unit/profiler/pandas/test_profiler_interface.py +++ b/ingestion/tests/unit/profiler/pandas/test_profiler_interface.py @@ -14,7 +14,7 @@ """ import os -from datetime import datetime, timezone +from datetime import datetime from unittest import TestCase, mock from uuid import uuid4 @@ -254,7 +254,7 @@ def test_get_all_metrics(self): table_profile = TableProfile( columnCount=profile_results["table"].get("columnCount"), rowCount=profile_results["table"].get(RowCount.name()), - timestamp=Timestamp(int(datetime.now(tz=timezone.utc).timestamp())), + timestamp=Timestamp(int(datetime.now().timestamp())), ) profile_request = CreateTableProfileRequest( diff --git a/ingestion/tests/unit/profiler/sqlalchemy/bigquery/test_map_struct.py b/ingestion/tests/unit/profiler/sqlalchemy/bigquery/test_map_struct.py index 87d78861eb66..92841ded5ff1 100644 --- a/ingestion/tests/unit/profiler/sqlalchemy/bigquery/test_map_struct.py +++ b/ingestion/tests/unit/profiler/sqlalchemy/bigquery/test_map_struct.py @@ -19,7 +19,7 @@ def test_map_struct(): - column = Column.parse_obj( + column = Column.model_validate( { "name": "col", "dataType": "STRUCT", diff --git a/ingestion/tests/unit/profiler/sqlalchemy/test_sqa_profiler_interface.py b/ingestion/tests/unit/profiler/sqlalchemy/test_sqa_profiler_interface.py index 95a4bd6c5785..7b78a6175a60 100644 --- a/ingestion/tests/unit/profiler/sqlalchemy/test_sqa_profiler_interface.py +++ b/ingestion/tests/unit/profiler/sqlalchemy/test_sqa_profiler_interface.py @@ -14,7 +14,7 @@ """ import os -from datetime import datetime, timezone +from datetime import datetime from unittest import TestCase from unittest.mock import patch from uuid import uuid4 @@ -228,7 +228,7 @@ def test_get_all_metrics(self): table_profile = TableProfile( columnCount=profile_results["table"].get("columnCount"), rowCount=profile_results["table"].get(RowCount.name()), - timestamp=Timestamp(int(datetime.now(tz=timezone.utc).timestamp())), + timestamp=Timestamp(int(datetime.now().timestamp())), ) profile_request = CreateTableProfileRequest( diff --git a/ingestion/tests/unit/test_databricks_lineage.py b/ingestion/tests/unit/test_databricks_lineage.py index 046ed08bb0d7..e6a5849da954 100644 --- a/ingestion/tests/unit/test_databricks_lineage.py +++ b/ingestion/tests/unit/test_databricks_lineage.py @@ -13,7 +13,7 @@ """ import json -from datetime import datetime, timezone +from datetime import datetime from pathlib import Path from unittest import TestCase from unittest.mock import patch @@ -38,7 +38,7 @@ userName="vijay@getcollate.io", startTime="1665566128192", endTime="1665566128329", - analysisDate=DateTime(datetime.now(tz=timezone.utc)), + analysisDate=DateTime(datetime.now()), aborted=None, serviceName="local_databricks1", databaseSchema=None, @@ -48,7 +48,7 @@ userName="vijay@getcollate.io", startTime="1665566127416", endTime="1665566127568", - analysisDate=DateTime(datetime.now(tz=timezone.utc)), + analysisDate=DateTime(datetime.now()), aborted=None, serviceName="local_databricks1", databaseSchema=None, @@ -58,7 +58,7 @@ userName="vijay@getcollate.io", startTime="1665566125414", endTime="1665566125579", - analysisDate=DateTime(datetime.now(tz=timezone.utc)), + analysisDate=DateTime(datetime.now()), aborted=None, serviceName="local_databricks1", databaseSchema=None, @@ -68,7 +68,7 @@ userName="vijay@getcollate.io", startTime="1665566124428", endTime="1665566124730", - analysisDate=DateTime(datetime.now(tz=timezone.utc)), + analysisDate=DateTime(datetime.now()), aborted=None, serviceName="local_databricks1", databaseSchema=None, @@ -122,7 +122,7 @@ class DatabricksLineageTests(TestCase): def __init__(self, methodName) -> None: super().__init__(methodName) - config = OpenMetadataWorkflowConfig.parse_obj(mock_databricks_config) + config = OpenMetadataWorkflowConfig.model_validate(mock_databricks_config) self.databricks = DatabricksLineageSource.create( mock_databricks_config["source"], diff --git a/ingestion/tests/unit/test_dbt.py b/ingestion/tests/unit/test_dbt.py index 5763c9505633..ba149c217008 100644 --- a/ingestion/tests/unit/test_dbt.py +++ b/ingestion/tests/unit/test_dbt.py @@ -262,7 +262,7 @@ class DbtUnitTest(TestCase): def __init__(self, methodName, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_dbt_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_dbt_config) self.dbt_source_obj = DbtSource.create( mock_dbt_config["source"], OpenMetadata(self.config.workflowConfig.openMetadataServerConfig), diff --git a/ingestion/tests/unit/test_handle_partitions.py b/ingestion/tests/unit/test_handle_partitions.py index 07fe9b73ce80..d88027498e39 100644 --- a/ingestion/tests/unit/test_handle_partitions.py +++ b/ingestion/tests/unit/test_handle_partitions.py @@ -118,7 +118,7 @@ def __init__( create_generic_connection.return_value = Mock() set_project_id.return_value = Mock() test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_bigquery_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_bigquery_config) self.bigquery_source = BigquerySource.create( mock_bigquery_config["source"], self.config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/test_ometa_to_dataframe.py b/ingestion/tests/unit/test_ometa_to_dataframe.py index 020c2cde97eb..53cb418270d5 100644 --- a/ingestion/tests/unit/test_ometa_to_dataframe.py +++ b/ingestion/tests/unit/test_ometa_to_dataframe.py @@ -54,7 +54,7 @@ def test_return_ometa_dataframes_sampled(self, test_connection): "metadata.mixins.pandas.pandas_mixin.fetch_dataframe", return_value=[resp_parquet_file], ): - config = OpenMetadataWorkflowConfig.parse_obj(mock_datalake_config) + config = OpenMetadataWorkflowConfig.model_validate(mock_datalake_config) datalake_source = DatalakeSource.create( mock_datalake_config["source"], config.workflowConfig.openMetadataServerConfig, @@ -88,7 +88,7 @@ def test_return_ometa_dataframes_sampled_fail(self, test_connection): return_value=None, ): with self.assertRaises(TypeError) as context: - config = OpenMetadataWorkflowConfig.parse_obj(mock_datalake_config) + config = OpenMetadataWorkflowConfig.model_validate(mock_datalake_config) datalake_source = DatalakeSource.create( mock_datalake_config["source"], config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/test_pgspider_lineage_unit.py b/ingestion/tests/unit/test_pgspider_lineage_unit.py index 52a275103058..a9a25bc0c5e0 100644 --- a/ingestion/tests/unit/test_pgspider_lineage_unit.py +++ b/ingestion/tests/unit/test_pgspider_lineage_unit.py @@ -586,7 +586,7 @@ class PGSpiderLineageUnitTests(TestCase): def __init__(self, methodName) -> None: super().__init__(methodName) - config = OpenMetadataWorkflowConfig.parse_obj(mock_pgspider_config) + config = OpenMetadataWorkflowConfig.model_validate(mock_pgspider_config) self.postgres = PostgresLineageSource.create( mock_pgspider_config["source"], config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/test_pydantic_v2.py b/ingestion/tests/unit/test_pydantic_v2.py index 32fa913b4f91..cea3b9028c18 100644 --- a/ingestion/tests/unit/test_pydantic_v2.py +++ b/ingestion/tests/unit/test_pydantic_v2.py @@ -9,7 +9,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """Test pydantic v2 models serialize data as pydantic v1""" -from datetime import datetime, timezone +from datetime import datetime from pydantic import AnyUrl from pydantic.v1 import BaseModel as BaseModelV1 @@ -78,7 +78,7 @@ class ModelV2(BaseModel): def test_tz_aware_date(): """Validate how we can create "aware" datetime objects""" - DateTime(datetime.now(tz=timezone.utc)) + DateTime(datetime.now()) def test_any_url(): diff --git a/ingestion/tests/unit/test_source_parsing.py b/ingestion/tests/unit/test_source_parsing.py index c54aed26629a..6a33147c0aca 100644 --- a/ingestion/tests/unit/test_source_parsing.py +++ b/ingestion/tests/unit/test_source_parsing.py @@ -124,7 +124,7 @@ def test_amundsen(): "sourceConfig": {"config": {"type": "DatabaseMetadata"}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, AmundsenConnection) @@ -164,7 +164,7 @@ def test_azure_sql(): }, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, AzureSQLConnection) @@ -182,7 +182,7 @@ def test_bigquery(): "privateKeyId": "privateKeyId", "privateKey": "privateKey", "clientEmail": "clientEmail", - "clientId": "clientId", + "clientId": "1234", "authUri": "https://accounts.google.com/o/oauth2/auth", "tokenUri": "https://oauth2.googleapis.com/token", "authProviderX509CertUrl": "https://www.googleapis.com/oauth2/v1/certs", @@ -194,7 +194,7 @@ def test_bigquery(): "sourceConfig": {"config": {"type": "DatabaseMetadata"}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, BigQueryConnection) @@ -225,7 +225,7 @@ def test_clickhouse(): }, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, ClickhouseConnection) @@ -245,7 +245,7 @@ def test_databricks(): "sourceConfig": {"config": {"type": "DatabaseMetadata"}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, DatabricksConnection) @@ -265,7 +265,7 @@ def test_db2(): "sourceConfig": {"config": {"type": "DatabaseMetadata"}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, Db2Connection) @@ -284,7 +284,7 @@ def test_deltalake(): "sourceConfig": {"config": {"type": "DatabaseMetadata"}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, DeltaLakeConnection) source = { @@ -299,7 +299,7 @@ def test_deltalake(): "sourceConfig": {"config": {"type": "DatabaseMetadata"}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, DeltaLakeConnection) source = { @@ -314,7 +314,7 @@ def test_deltalake(): "sourceConfig": {"config": {"type": "DatabaseMetadata"}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, DeltaLakeConnection) @@ -346,7 +346,7 @@ def test_dynamo_db(): }, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, DynamoDBConnection) @@ -368,7 +368,7 @@ def test_glue(): "sourceConfig": {"config": {"type": "DatabaseMetadata"}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, GlueConnection) @@ -382,7 +382,7 @@ def test_hive(): "sourceConfig": {"config": {"type": "DatabaseMetadata"}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, HiveConnection) @@ -396,7 +396,7 @@ def test_impala(): "sourceConfig": {"config": {"type": "DatabaseMetadata"}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, ImpalaConnection) @@ -423,7 +423,7 @@ def test_looker(): "sourceConfig": {"config": {}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, LookerConnection) @@ -442,7 +442,7 @@ def test_mariadb(): "sourceConfig": {"config": {"type": "DatabaseMetadata"}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, MariaDBConnection) @@ -461,7 +461,7 @@ def test_mariadb(): "sourceConfig": {"config": {"type": "DatabaseMetadata"}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, MariaDBConnection) @@ -482,7 +482,7 @@ def test_metabase(): }, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, MetabaseConnection) @@ -510,7 +510,7 @@ def test_mssql(): "sourceConfig": {"config": {"type": "DatabaseMetadata"}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, MssqlConnection) @@ -529,7 +529,7 @@ def test_mysql(): "sourceConfig": {"config": {"type": "DatabaseMetadata"}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, MysqlConnection) @@ -549,7 +549,7 @@ def test_oracle(): "sourceConfig": {"config": {}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, OracleConnection) @@ -571,7 +571,7 @@ def test_postgres(): "sourceConfig": {"config": {"type": "DatabaseMetadata"}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, PostgresConnection) @@ -597,7 +597,7 @@ def test_powerbi(): }, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, PowerBIConnection) @@ -617,7 +617,7 @@ def test_presto(): "sourceConfig": {"config": {"generateSampleData": False}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, PrestoConnection) @@ -638,7 +638,7 @@ def test_redash(): }, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, RedashConnection) @@ -664,7 +664,7 @@ def test_redshift(): }, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, RedshiftConnection) @@ -689,7 +689,7 @@ def test_salesforce(): "sourceConfig": {"config": {"type": "DatabaseMetadata"}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, SalesforceConnection) @@ -707,7 +707,7 @@ def test_sample_data(): "sourceConfig": {}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance( config.serviceConnection.root.config, customDatabaseConnection.CustomDatabaseConnection, @@ -729,7 +729,7 @@ def test_singlestore(): "sourceConfig": {"config": {"type": "DatabaseMetadata"}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, SingleStoreConnection) @@ -762,7 +762,7 @@ def test_snowflake(): }, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, SnowflakeConnection) @@ -774,7 +774,7 @@ def test_sqlite(): "sourceConfig": {"config": {}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, SQLiteConnection) @@ -798,7 +798,7 @@ def test_superset(): }, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, SupersetConnection) @@ -821,7 +821,7 @@ def test_tableau(): }, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, TableauConnection) @@ -840,7 +840,7 @@ def test_trino(): "sourceConfig": {"config": {}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, TrinoConnection) @@ -860,5 +860,5 @@ def test_vertica(): "sourceConfig": {"config": {"type": "DatabaseMetadata"}}, } - config: WorkflowSource = WorkflowSource.parse_obj(source) + config: WorkflowSource = WorkflowSource.model_validate(source) assert isinstance(config.serviceConnection.root.config, VerticaConnection) diff --git a/ingestion/tests/unit/test_usage_filter.py b/ingestion/tests/unit/test_usage_filter.py index 45ad3b040868..d44286d6ab7a 100644 --- a/ingestion/tests/unit/test_usage_filter.py +++ b/ingestion/tests/unit/test_usage_filter.py @@ -146,7 +146,7 @@ class UsageQueryFilterTests(TestCase): @patch.object(OpenMetadata, "list_all_entities", mock_list_entities) def test_prepare_clickhouse(self): - config = OpenMetadataWorkflowConfig.parse_obj(mock_clickhouse_config) + config = OpenMetadataWorkflowConfig.model_validate(mock_clickhouse_config) clickhouse_source = ClickhouseUsageSource.create( mock_clickhouse_config["source"], OpenMetadata(config.workflowConfig.openMetadataServerConfig), diff --git a/ingestion/tests/unit/test_usage_log.py b/ingestion/tests/unit/test_usage_log.py index ccb1b4b7f41f..5b4bad3eba66 100644 --- a/ingestion/tests/unit/test_usage_log.py +++ b/ingestion/tests/unit/test_usage_log.py @@ -120,7 +120,7 @@ class QueryLogSourceTest(TestCase): def __init__(self, methodName) -> None: super().__init__(methodName) - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_query_log_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_query_log_config) self.source = QueryLogUsageSource.create( mock_query_log_config["source"], self.config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/dashboard/test_domodashboard.py b/ingestion/tests/unit/topology/dashboard/test_domodashboard.py index c1caa71f7de1..66b35371184a 100644 --- a/ingestion/tests/unit/topology/dashboard/test_domodashboard.py +++ b/ingestion/tests/unit/topology/dashboard/test_domodashboard.py @@ -150,7 +150,9 @@ def __init__(self, methodName, domo_client, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False domo_client.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_domopipeline_config) + self.config = OpenMetadataWorkflowConfig.model_validate( + mock_domopipeline_config + ) self.domodashboard = DomodashboardSource.create( mock_domopipeline_config["source"], self.config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/dashboard/test_looker.py b/ingestion/tests/unit/topology/dashboard/test_looker.py index 83cebd3df867..dfabd029c768 100644 --- a/ingestion/tests/unit/topology/dashboard/test_looker.py +++ b/ingestion/tests/unit/topology/dashboard/test_looker.py @@ -136,7 +136,7 @@ class LookerUnitTest(TestCase): def __init__(self, methodName, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(MOCK_LOOKER_CONFIG) + self.config = OpenMetadataWorkflowConfig.model_validate(MOCK_LOOKER_CONFIG) # This already validates that the source can be initialized self.looker: LookerSource = LookerSource.create( diff --git a/ingestion/tests/unit/topology/dashboard/test_metabase.py b/ingestion/tests/unit/topology/dashboard/test_metabase.py index 3ae3a6177257..4aafaa09575e 100644 --- a/ingestion/tests/unit/topology/dashboard/test_metabase.py +++ b/ingestion/tests/unit/topology/dashboard/test_metabase.py @@ -234,7 +234,7 @@ def __init__(self, methodName, get_connection, test_connection) -> None: super().__init__(methodName) get_connection.return_value = False test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_config) self.metabase: MetabaseSource = MetabaseSource.create( mock_config["source"], OpenMetadata(self.config.workflowConfig.openMetadataServerConfig), diff --git a/ingestion/tests/unit/topology/dashboard/test_qlikcloud.py b/ingestion/tests/unit/topology/dashboard/test_qlikcloud.py index 4596ba01c653..5c72ae07c131 100644 --- a/ingestion/tests/unit/topology/dashboard/test_qlikcloud.py +++ b/ingestion/tests/unit/topology/dashboard/test_qlikcloud.py @@ -160,7 +160,9 @@ def __init__(self, methodName) -> None: with patch.object(QlikCloudClient, "get_dashboards_list", return_value=None): super().__init__(methodName) # test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_qlikcloud_config) + self.config = OpenMetadataWorkflowConfig.model_validate( + mock_qlikcloud_config + ) self.qlikcloud = QlikcloudSource.create( mock_qlikcloud_config["source"], OpenMetadata(self.config.workflowConfig.openMetadataServerConfig), diff --git a/ingestion/tests/unit/topology/dashboard/test_qliksense.py b/ingestion/tests/unit/topology/dashboard/test_qliksense.py index 44d02ca7b158..6230e4c2f11e 100644 --- a/ingestion/tests/unit/topology/dashboard/test_qliksense.py +++ b/ingestion/tests/unit/topology/dashboard/test_qliksense.py @@ -172,7 +172,9 @@ def __init__(self, methodName) -> None: ): super().__init__(methodName) # test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_qliksense_config) + self.config = OpenMetadataWorkflowConfig.model_validate( + mock_qliksense_config + ) self.qliksense = QliksenseSource.create( mock_qliksense_config["source"], OpenMetadata(self.config.workflowConfig.openMetadataServerConfig), diff --git a/ingestion/tests/unit/topology/dashboard/test_quicksight.py b/ingestion/tests/unit/topology/dashboard/test_quicksight.py index d94ae1477d6a..7aa9208dee6c 100644 --- a/ingestion/tests/unit/topology/dashboard/test_quicksight.py +++ b/ingestion/tests/unit/topology/dashboard/test_quicksight.py @@ -154,7 +154,7 @@ class QuickSightUnitTest(TestCase): def __init__(self, methodName, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_quicksight_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_quicksight_config) self.quicksight = QuicksightSource.create( mock_quicksight_config["source"], self.config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/dashboard/test_superset.py b/ingestion/tests/unit/topology/dashboard/test_superset.py index cca15fd59e13..39db3f5812fb 100644 --- a/ingestion/tests/unit/topology/dashboard/test_superset.py +++ b/ingestion/tests/unit/topology/dashboard/test_superset.py @@ -403,7 +403,9 @@ def __init__(self, methodName) -> None: }, }, } - self.config = OpenMetadataWorkflowConfig.parse_obj(MOCK_SUPERSET_API_CONFIG) + self.config = OpenMetadataWorkflowConfig.model_validate( + MOCK_SUPERSET_API_CONFIG + ) self.superset_api: SupersetSource = SupersetSource.create( MOCK_SUPERSET_API_CONFIG["source"], diff --git a/ingestion/tests/unit/topology/dashboard/test_tableau.py b/ingestion/tests/unit/topology/dashboard/test_tableau.py index 234ba806ff2f..1fe1b5fa183e 100644 --- a/ingestion/tests/unit/topology/dashboard/test_tableau.py +++ b/ingestion/tests/unit/topology/dashboard/test_tableau.py @@ -176,7 +176,7 @@ def __init__( get_connection.return_value = False tableau_server_connection.return_value = False test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_tableau_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_tableau_config) self.tableau = TableauSource.create( mock_tableau_config["source"], self.config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/database/test_bigquery.py b/ingestion/tests/unit/topology/database/test_bigquery.py index db0e3e99d200..1fb1b7d33130 100644 --- a/ingestion/tests/unit/topology/database/test_bigquery.py +++ b/ingestion/tests/unit/topology/database/test_bigquery.py @@ -71,7 +71,7 @@ # this is a valid key that was generated on a local machine and is not used for any real project "privateKey": "-----BEGIN RSA PRIVATE KEY-----\nMIIEpQIBAAKCAQEAw3vHG9fDIkcYB0xi2Mv4fS2gUzKR9ZRrcVNeKkqGFTT71AVB\nOzgIqYVe8b2aWODuNye6sipcrqTqOt05Esj+sxhk5McM9bE2RlxXC5QH/Bp9zxMP\n/Yksv9Ov7fdDt/loUk7sTXvI+7LDJfmRYU6MtVjyyLs7KpQIB2xBWEToU1xZY+v0\ndRC1NA+YWc+FjXbAiFAf9d4gXkYO8VmU5meixVh4C8nsjokEXk0T/HEItpZCxadk\ndZ7LKUE/HDmWCO2oNG6sCf4ET2crjSdYIfXuREopX1aQwnk7KbI4/YIdlRz1I369\nAz3+Hxlf9lLJVH3+itN4GXrR9yWWKWKDnwDPbQIDAQABAoIBAQC3X5QuTR7SN8iV\niBUtc2D84+ECSmza5shG/UJW/6N5n0Mf53ICgBS4GNEwiYCRISa0/ILIgK6CcVb7\nsuvH8F3kWNzEMui4TO0x4YsR5GH9HkioCCS224frxkLBQnL20HIIy9ok8Rpe6Zjg\nNZUnp4yczPyqSeA9l7FUbTt69uDM2Cx61m8REOpFukpnYLyZGbmNPYmikEO+rq9r\nwNID5dkSeVuQYo4MQdRavOGFUWvUYXzkEQ0A6vPyraVBfolESX8WaLNVjic7nIa3\nujdSNojnJqGJ3gslntcmN1d4JOfydc4bja4/NdNlcOHpWDGLzY1QnaDe0Koxn8sx\nLT9MVD2NAoGBAPy7r726bKVGWcwqTzUuq1OWh5c9CAc4N2zWBBldSJyUdllUq52L\nWTyva6GRoRzCcYa/dKLLSM/k4eLf9tpxeIIfTOMsvzGtbAdm257ndMXNvfYpxCfU\nK/gUFfAUGHZ3MucTHRY6DTkJg763Sf6PubA2fqv3HhVZDK/1HGDtHlTPAoGBAMYC\npdV7O7lAyXS/d9X4PQZ4BM+P8MbXEdGBbPPlzJ2YIb53TEmYfSj3z41u9+BNnhGP\n4uzUyAR/E4sxrA2+Ll1lPSCn+KY14WWiVGfWmC5j1ftdpkbrXstLN8NpNYzrKZwx\njdR0ZkwvZ8B5+kJ1hK96giwWS+SJxJR3TohcQ18DAoGAJSfmv2r//BBqtURnHrd8\nwq43wvlbC8ytAVg5hA0d1r9Q4vM6w8+vz+cuWLOTTyobDKdrG1/tlXrd5r/sh9L0\n15SIdkGm3kPTxQbPNP5sQYRs8BrV1tEvoao6S3B45DnEBwrdVN42AXOvpcNGoqE4\nuHpahyeuiY7s+ZV8lZdmxSsCgYEAolr5bpmk1rjwdfGoaKEqKGuwRiBX5DHkQkxE\n8Zayt2VOBcX7nzyRI05NuEIMrLX3rZ61CktN1aH8fF02He6aRaoE/Qm9L0tujM8V\nNi8WiLMDeR/Ifs3u4/HAv1E8v1byv0dCa7klR8J257McJ/ID4X4pzcxaXgE4ViOd\nGOHNu9ECgYEApq1zkZthEQymTUxs+lSFcubQpaXyf5ZC61cJewpWkqGDtSC+8DxE\nF/jydybWuoNHXymnvY6QywxuIooivbuib6AlgpEJeybmnWlDOZklFOD0abNZ+aNO\ndUk7XVGffCakXQ0jp1kmZA4lGsYK1h5dEU5DgXqu4UYJ88Vttax2W+Y=\n-----END RSA PRIVATE KEY-----\n", "clientEmail": "gcpuser@project_id.iam.gserviceaccount.com", - "clientId": "client_id", + "clientId": "1234", "authUri": "https://accounts.google.com/o/oauth2/auth", "tokenUri": "https://oauth2.googleapis.com/token", "authProviderX509CertUrl": "https://www.googleapis.com/oauth2/v1/certs", @@ -697,7 +697,7 @@ def __init__( "credentials" ]["gcpConfig"] = "credentials.json" - self.config = OpenMetadataWorkflowConfig.parse_obj( + self.config = OpenMetadataWorkflowConfig.model_validate( mock_credentials_path_bq_config ) self.bq_query_parser = BigqueryLineageSource( diff --git a/ingestion/tests/unit/topology/database/test_couchbase.py b/ingestion/tests/unit/topology/database/test_couchbase.py index 649f75f59cbb..8ac011fbb234 100644 --- a/ingestion/tests/unit/topology/database/test_couchbase.py +++ b/ingestion/tests/unit/topology/database/test_couchbase.py @@ -188,7 +188,7 @@ def __init__(self, methodName, get_connection, test_connection) -> None: get_connection.return_value = False test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_couch_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_couch_config) self.couch_source = CouchbaseSource.create( mock_couch_config["source"], OpenMetadata(self.config.workflowConfig.openMetadataServerConfig), diff --git a/ingestion/tests/unit/topology/database/test_databricks.py b/ingestion/tests/unit/topology/database/test_databricks.py index 462993f593c3..b858bc9a1818 100644 --- a/ingestion/tests/unit/topology/database/test_databricks.py +++ b/ingestion/tests/unit/topology/database/test_databricks.py @@ -262,7 +262,7 @@ def __init__(self, methodName, test_connection, db_init_version) -> None: test_connection.return_value = False db_init_version.return_value = None - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_databricks_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_databricks_config) self.databricks_source = DatabricksSource.create( mock_databricks_config["source"], self.config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/database/test_datalake.py b/ingestion/tests/unit/topology/database/test_datalake.py index 6acf57180456..847fb9c7cda3 100644 --- a/ingestion/tests/unit/topology/database/test_datalake.py +++ b/ingestion/tests/unit/topology/database/test_datalake.py @@ -471,7 +471,7 @@ class DatalakeUnitTest(TestCase): def __init__(self, methodName, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_datalake_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_datalake_config) self.datalake_source = DatalakeSource.create( mock_datalake_config["source"], self.config.workflowConfig.openMetadataServerConfig, @@ -563,7 +563,7 @@ def test_avro_file_parse(self): "privateKeyId": "private_key_id", "privateKey": "private_key", "clientEmail": "gcpuser@project_id.iam.gserviceaccount.com", - "clientId": "client_id", + "clientId": "1234", "authUri": "https://accounts.google.com/o/oauth2/auth", "tokenUri": "https://oauth2.googleapis.com/token", "authProviderX509CertUrl": "https://www.googleapis.com/oauth2/v1/certs", @@ -610,7 +610,9 @@ class DatalakeGCSUnitTest(TestCase): def __init__(self, methodName, _, __, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_datalake_gcs_config) + self.config = OpenMetadataWorkflowConfig.model_validate( + mock_datalake_gcs_config + ) self.datalake_source = DatalakeSource.create( mock_datalake_gcs_config["source"], self.config.workflowConfig.openMetadataServerConfig, @@ -632,7 +634,7 @@ def test_multiple_project_id_implementation( ): self.datalake_source_multiple_project_id = DatalakeSource.create( mock_multiple_project_id["source"], - OpenMetadataWorkflowConfig.parse_obj( + OpenMetadataWorkflowConfig.model_validate( mock_multiple_project_id ).workflowConfig.openMetadataServerConfig, ) diff --git a/ingestion/tests/unit/topology/database/test_deltalake.py b/ingestion/tests/unit/topology/database/test_deltalake.py index 1e15a1f67529..d44844547b9b 100644 --- a/ingestion/tests/unit/topology/database/test_deltalake.py +++ b/ingestion/tests/unit/topology/database/test_deltalake.py @@ -111,7 +111,7 @@ class DeltaLakeUnitTest(TestCase): Add method validations from Deltalake ingestion """ - config: OpenMetadataWorkflowConfig = OpenMetadataWorkflowConfig.parse_obj( + config: OpenMetadataWorkflowConfig = OpenMetadataWorkflowConfig.model_validate( MOCK_DELTA_CONFIG ) delta: DeltalakeSource = DeltalakeSource.create( diff --git a/ingestion/tests/unit/topology/database/test_domodatabase.py b/ingestion/tests/unit/topology/database/test_domodatabase.py index 177d28092575..cda6ee42ab52 100644 --- a/ingestion/tests/unit/topology/database/test_domodatabase.py +++ b/ingestion/tests/unit/topology/database/test_domodatabase.py @@ -254,7 +254,9 @@ def __init__( ) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_domodatabase_config) + self.config = OpenMetadataWorkflowConfig.model_validate( + mock_domodatabase_config + ) self.domodatabase = DomodatabaseSource.create( mock_domodatabase_config["source"], self.config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/database/test_doris.py b/ingestion/tests/unit/topology/database/test_doris.py index 9998821c8d58..16e856461c70 100644 --- a/ingestion/tests/unit/topology/database/test_doris.py +++ b/ingestion/tests/unit/topology/database/test_doris.py @@ -66,7 +66,7 @@ class DorisUnitTest(TestCase): def __init__(self, methodName, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_doris_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_doris_config) self.doris_source = DorisSource.create( mock_doris_config["source"], self.config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/database/test_glue.py b/ingestion/tests/unit/topology/database/test_glue.py index 850881332068..36f88a8e6f49 100644 --- a/ingestion/tests/unit/topology/database/test_glue.py +++ b/ingestion/tests/unit/topology/database/test_glue.py @@ -132,7 +132,7 @@ class GlueUnitTest(TestCase): def __init__(self, methodName, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_glue_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_glue_config) self.glue_source = GlueSource.create( mock_glue_config["source"], self.config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/database/test_greenplum.py b/ingestion/tests/unit/topology/database/test_greenplum.py index d423e64c0874..40f45db44030 100644 --- a/ingestion/tests/unit/topology/database/test_greenplum.py +++ b/ingestion/tests/unit/topology/database/test_greenplum.py @@ -65,7 +65,7 @@ class greenplumUnitTest(TestCase): def __init__(self, methodName, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_greenplum_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_greenplum_config) self.greenplum_source = GreenplumSource.create( mock_greenplum_config["source"], self.config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/database/test_hive.py b/ingestion/tests/unit/topology/database/test_hive.py index 744400edb22a..f852585bb872 100644 --- a/ingestion/tests/unit/topology/database/test_hive.py +++ b/ingestion/tests/unit/topology/database/test_hive.py @@ -331,7 +331,7 @@ def __init__( ) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_hive_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_hive_config) self.hive = HiveSource.create( mock_hive_config["source"], self.config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/database/test_iceberg.py b/ingestion/tests/unit/topology/database/test_iceberg.py index 5a1cc44806e1..e46c633cb92f 100644 --- a/ingestion/tests/unit/topology/database/test_iceberg.py +++ b/ingestion/tests/unit/topology/database/test_iceberg.py @@ -691,7 +691,7 @@ def test_get_owner_ref(self): self.iceberg.context.get().database_schema, table_name, ), - "metadata": TableMetadataV2.parse_obj( + "metadata": TableMetadataV2.model_validate( { "location": "foo", "last_column_id": 1, @@ -743,7 +743,7 @@ def test_get_owner_ref(self): self.iceberg.context.get().database_schema, table_name, ), - "metadata": TableMetadataV2.parse_obj( + "metadata": TableMetadataV2.model_validate( { "location": "foo", "last_column_id": 1, @@ -794,7 +794,7 @@ def test_yield_table(self): self.iceberg.context.get().database_schema, table_name, ), - "metadata": TableMetadataV2.parse_obj( + "metadata": TableMetadataV2.model_validate( { "location": "foo", "current-schema-id": 0, diff --git a/ingestion/tests/unit/topology/database/test_mongodb.py b/ingestion/tests/unit/topology/database/test_mongodb.py index 8a8fca4934bb..8004a2804f61 100644 --- a/ingestion/tests/unit/topology/database/test_mongodb.py +++ b/ingestion/tests/unit/topology/database/test_mongodb.py @@ -195,7 +195,7 @@ class MongoDBUnitTest(TestCase): def __init__(self, methodName, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_mongo_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_mongo_config) self.mongo_source = MongodbSource.create( mock_mongo_config["source"], OpenMetadata(self.config.workflowConfig.openMetadataServerConfig), diff --git a/ingestion/tests/unit/topology/database/test_mssql.py b/ingestion/tests/unit/topology/database/test_mssql.py index 5c8d073824a0..16680243dece 100644 --- a/ingestion/tests/unit/topology/database/test_mssql.py +++ b/ingestion/tests/unit/topology/database/test_mssql.py @@ -303,7 +303,7 @@ def __init__( ) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_mssql_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_mssql_config) self.mssql = MssqlSource.create( mock_mssql_config["source"], self.config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/database/test_mysql.py b/ingestion/tests/unit/topology/database/test_mysql.py index 97ed3f056484..d69a2f595770 100644 --- a/ingestion/tests/unit/topology/database/test_mysql.py +++ b/ingestion/tests/unit/topology/database/test_mysql.py @@ -66,7 +66,7 @@ class MysqlUnitTest(TestCase): def __init__(self, methodName, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_mysql_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_mysql_config) self.mysql_source = MysqlSource.create( mock_mysql_config["source"], self.config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/database/test_oracle.py b/ingestion/tests/unit/topology/database/test_oracle.py index 999bfabe4e28..a4fe26a2174d 100644 --- a/ingestion/tests/unit/topology/database/test_oracle.py +++ b/ingestion/tests/unit/topology/database/test_oracle.py @@ -183,9 +183,9 @@ def __init__( ) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_oracle_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_oracle_config) self.metadata = OpenMetadata( - OpenMetadataConnection.parse_obj( + OpenMetadataConnection.model_validate( mock_oracle_config["workflowConfig"]["openMetadataServerConfig"] ) ) diff --git a/ingestion/tests/unit/topology/database/test_postgres.py b/ingestion/tests/unit/topology/database/test_postgres.py index 8adc0da9bcdc..86da89043473 100644 --- a/ingestion/tests/unit/topology/database/test_postgres.py +++ b/ingestion/tests/unit/topology/database/test_postgres.py @@ -280,7 +280,7 @@ class PostgresUnitTest(TestCase): def __init__(self, methodName, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_postgres_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_postgres_config) self.postgres_source = PostgresSource.create( mock_postgres_config["source"], self.config.workflowConfig.openMetadataServerConfig, @@ -296,7 +296,7 @@ def __init__(self, methodName, test_connection) -> None: "database_schema" ] = MOCK_DATABASE_SCHEMA.name.root - self.usage_config = OpenMetadataWorkflowConfig.parse_obj( + self.usage_config = OpenMetadataWorkflowConfig.model_validate( mock_postgres_usage_config ) self.postgres_usage_source = PostgresUsageSource.create( diff --git a/ingestion/tests/unit/topology/database/test_redshift.py b/ingestion/tests/unit/topology/database/test_redshift.py index f10dc8189c18..04095e6d2ac9 100644 --- a/ingestion/tests/unit/topology/database/test_redshift.py +++ b/ingestion/tests/unit/topology/database/test_redshift.py @@ -65,7 +65,7 @@ class RedshiftUnitTest(TestCase): def __init__(self, methodName, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_redshift_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_redshift_config) self.redshift_source = RedshiftSource.create( mock_redshift_config["source"], self.config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/database/test_salesforce.py b/ingestion/tests/unit/topology/database/test_salesforce.py index fb57b68e826d..c2d8a866fb33 100644 --- a/ingestion/tests/unit/topology/database/test_salesforce.py +++ b/ingestion/tests/unit/topology/database/test_salesforce.py @@ -437,7 +437,7 @@ class SalesforceUnitTest(TestCase): def __init__(self, methodName, salesforce, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_salesforce_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_salesforce_config) self.salesforce_source = SalesforceSource.create( mock_salesforce_config["source"], self.config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/database/test_sas.py b/ingestion/tests/unit/topology/database/test_sas.py index 19850dfe1f9a..ba85bfa3a3e3 100644 --- a/ingestion/tests/unit/topology/database/test_sas.py +++ b/ingestion/tests/unit/topology/database/test_sas.py @@ -252,13 +252,13 @@ class SASUnitTest(TestCase): def __init__(self, method_name, test_connection) -> None: super().__init__(method_name) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_sas_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_sas_config) self.sas_source = SasSource.create( mock_sas_config["source"], OpenMetadata(self.config.workflowConfig.openMetadataServerConfig), ) self.metadata = OpenMetadata( - OpenMetadataConnection.parse_obj( + OpenMetadataConnection.model_validate( mock_sas_config["workflowConfig"]["openMetadataServerConfig"] ) ) diff --git a/ingestion/tests/unit/topology/database/test_snowflake.py b/ingestion/tests/unit/topology/database/test_snowflake.py index 1fff76e4447a..14fde4908c46 100644 --- a/ingestion/tests/unit/topology/database/test_snowflake.py +++ b/ingestion/tests/unit/topology/database/test_snowflake.py @@ -128,7 +128,7 @@ def get_snowflake_sources(): "metadata.ingestion.source.database.common_db_source.CommonDbSourceService.test_connection", return_value=False, ): - config = OpenMetadataWorkflowConfig.parse_obj( + config = OpenMetadataWorkflowConfig.model_validate( SNOWFLAKE_CONFIGURATIONS["not_incremental"] ) sources["not_incremental"] = SnowflakeSource.create( @@ -141,7 +141,7 @@ def get_snowflake_sources(): "metadata.ingestion.source.database.incremental_metadata_extraction.IncrementalConfigCreator._get_pipeline_statuses", return_value=MOCK_PIPELINE_STATUSES, ): - config = OpenMetadataWorkflowConfig.parse_obj( + config = OpenMetadataWorkflowConfig.model_validate( SNOWFLAKE_CONFIGURATIONS["incremental"] ) sources["incremental"] = SnowflakeSource.create( diff --git a/ingestion/tests/unit/topology/metadata/test_amundsen.py b/ingestion/tests/unit/topology/metadata/test_amundsen.py index 1668ca82d91a..9d5b949cd6dd 100644 --- a/ingestion/tests/unit/topology/metadata/test_amundsen.py +++ b/ingestion/tests/unit/topology/metadata/test_amundsen.py @@ -175,7 +175,7 @@ class AmundsenUnitTest(TestCase): def __init__(self, methodName, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_amundsen_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_amundsen_config) self.amundsen = AmundsenSource.create( mock_amundsen_config["source"], OpenMetadata(self.config.workflowConfig.openMetadataServerConfig), diff --git a/ingestion/tests/unit/topology/metadata/test_atlas.py b/ingestion/tests/unit/topology/metadata/test_atlas.py index 49b3644b777c..c5b9a020c707 100644 --- a/ingestion/tests/unit/topology/metadata/test_atlas.py +++ b/ingestion/tests/unit/topology/metadata/test_atlas.py @@ -293,13 +293,13 @@ class AtlasUnitTest(TestCase): def __init__(self, methodName, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_atlas_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_atlas_config) self.atlas_source = AtlasSource.create( mock_atlas_config["source"], OpenMetadata(self.config.workflowConfig.openMetadataServerConfig), ) self.metadata = OpenMetadata( - OpenMetadataConnection.parse_obj( + OpenMetadataConnection.model_validate( mock_atlas_config["workflowConfig"]["openMetadataServerConfig"] ) ) diff --git a/ingestion/tests/unit/topology/pipeline/test_airbyte.py b/ingestion/tests/unit/topology/pipeline/test_airbyte.py index 2f7cdc095e10..63f6a64ae3ca 100644 --- a/ingestion/tests/unit/topology/pipeline/test_airbyte.py +++ b/ingestion/tests/unit/topology/pipeline/test_airbyte.py @@ -165,7 +165,7 @@ class AirbyteUnitTest(TestCase): def __init__(self, methodName, airbyte_client, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - config = OpenMetadataWorkflowConfig.parse_obj(mock_airbyte_config) + config = OpenMetadataWorkflowConfig.model_validate(mock_airbyte_config) self.airbyte = AirbyteSource.create( mock_airbyte_config["source"], config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/pipeline/test_airflow.py b/ingestion/tests/unit/topology/pipeline/test_airflow.py index 5a9a11dcb04e..2d2165df8933 100644 --- a/ingestion/tests/unit/topology/pipeline/test_airflow.py +++ b/ingestion/tests/unit/topology/pipeline/test_airflow.py @@ -168,7 +168,7 @@ class TestAirflow(TestCase): def __init__(self, methodName, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(MOCK_CONFIG) + self.config = OpenMetadataWorkflowConfig.model_validate(MOCK_CONFIG) # This already validates that the source can be initialized self.airflow: AirflowSource = AirflowSource.create( @@ -187,7 +187,7 @@ def test_parsing(self): dag = AirflowDagDetails( dag_id="id", fileloc="loc", - data=AirflowDag.parse_obj(SERIALIZED_DAG), + data=AirflowDag.model_validate(SERIALIZED_DAG), max_active_runs=data.get("max_active_runs", None), description=data.get("_description", None), start_date=data.get("start_date", None), diff --git a/ingestion/tests/unit/topology/pipeline/test_dagster.py b/ingestion/tests/unit/topology/pipeline/test_dagster.py index 402c04e4fac6..48dcceef2a9f 100644 --- a/ingestion/tests/unit/topology/pipeline/test_dagster.py +++ b/ingestion/tests/unit/topology/pipeline/test_dagster.py @@ -269,7 +269,7 @@ def __init__(self, methodName, graphql_client, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False graphql_client.return_value = False - config = OpenMetadataWorkflowConfig.parse_obj(mock_dagster_config) + config = OpenMetadataWorkflowConfig.model_validate(mock_dagster_config) self.dagster = DagsterSource.create( mock_dagster_config["source"], config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/pipeline/test_databricks_pipeline.py b/ingestion/tests/unit/topology/pipeline/test_databricks_pipeline.py index 3a248bf78144..67f1660502fa 100644 --- a/ingestion/tests/unit/topology/pipeline/test_databricks_pipeline.py +++ b/ingestion/tests/unit/topology/pipeline/test_databricks_pipeline.py @@ -230,7 +230,7 @@ def __init__(self, methodName, test_connection) -> None: super().__init__(methodName) log_ansi_encoded_string(message="init") test_connection.return_value = False - config = OpenMetadataWorkflowConfig.parse_obj(mock_databricks_config) + config = OpenMetadataWorkflowConfig.model_validate(mock_databricks_config) self.databricks = DatabrickspipelineSource.create( mock_databricks_config["source"], diff --git a/ingestion/tests/unit/topology/pipeline/test_domopipeline.py b/ingestion/tests/unit/topology/pipeline/test_domopipeline.py index 639364996158..b5f58b09453a 100644 --- a/ingestion/tests/unit/topology/pipeline/test_domopipeline.py +++ b/ingestion/tests/unit/topology/pipeline/test_domopipeline.py @@ -3,7 +3,7 @@ """ import json -from datetime import datetime, timezone +from datetime import datetime from pathlib import Path from unittest import TestCase from unittest.mock import patch @@ -151,7 +151,7 @@ sourceUrl=None, concurrency=None, pipelineLocation=None, - startDate=datetime(2022, 10, 7, 13, 20, 16, tzinfo=timezone.utc), + startDate=datetime(2022, 10, 7, 13, 20, 16), tasks=[ Task( name="1", @@ -247,7 +247,9 @@ def __init__(self, methodName, domo_client, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False domo_client.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_domopipeline_config) + self.config = OpenMetadataWorkflowConfig.model_validate( + mock_domopipeline_config + ) self.domopipeline = DomopipelineSource.create( mock_domopipeline_config["source"], self.config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/pipeline/test_fivetran.py b/ingestion/tests/unit/topology/pipeline/test_fivetran.py index f50ab07369ba..0261dcb7540b 100644 --- a/ingestion/tests/unit/topology/pipeline/test_fivetran.py +++ b/ingestion/tests/unit/topology/pipeline/test_fivetran.py @@ -120,7 +120,7 @@ class FivetranUnitTest(TestCase): def __init__(self, methodName, fivetran_client, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - config = OpenMetadataWorkflowConfig.parse_obj(mock_fivetran_config) + config = OpenMetadataWorkflowConfig.model_validate(mock_fivetran_config) self.fivetran = FivetranSource.create( mock_fivetran_config["source"], config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/pipeline/test_nifi.py b/ingestion/tests/unit/topology/pipeline/test_nifi.py index 2b339298913f..4dae439db98d 100644 --- a/ingestion/tests/unit/topology/pipeline/test_nifi.py +++ b/ingestion/tests/unit/topology/pipeline/test_nifi.py @@ -202,7 +202,7 @@ def __init__(self, methodName, nifi_token_prop, test_connection) -> None: nifi_token_prop.return_value.token.return_value = "token" - config = OpenMetadataWorkflowConfig.parse_obj(mock_nifi_config) + config = OpenMetadataWorkflowConfig.model_validate(mock_nifi_config) self.nifi = NifiSource.create( mock_nifi_config["source"], config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/pipeline/test_openlineage.py b/ingestion/tests/unit/topology/pipeline/test_openlineage.py index d0d60f3ef267..52bd633908a2 100644 --- a/ingestion/tests/unit/topology/pipeline/test_openlineage.py +++ b/ingestion/tests/unit/topology/pipeline/test_openlineage.py @@ -133,7 +133,7 @@ class OpenLineageUnitTest(unittest.TestCase): def __init__(self, methodName, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - config = OpenMetadataWorkflowConfig.parse_obj(MOCK_OL_CONFIG) + config = OpenMetadataWorkflowConfig.model_validate(MOCK_OL_CONFIG) self.open_lineage_source = OpenlineageSource.create( MOCK_OL_CONFIG["source"], config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/pipeline/test_spline.py b/ingestion/tests/unit/topology/pipeline/test_spline.py index 8c7ff4f75c61..9d72bb4b3453 100644 --- a/ingestion/tests/unit/topology/pipeline/test_spline.py +++ b/ingestion/tests/unit/topology/pipeline/test_spline.py @@ -239,7 +239,7 @@ class SplineUnitTest(TestCase): def __init__(self, methodName, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - config = OpenMetadataWorkflowConfig.parse_obj(mock_spline_config) + config = OpenMetadataWorkflowConfig.model_validate(mock_spline_config) self.spline = SplineSource.create( mock_spline_config["source"], config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/search/test_elasticsearch.py b/ingestion/tests/unit/topology/search/test_elasticsearch.py index a9dba0f99817..ec8109dc9e39 100644 --- a/ingestion/tests/unit/topology/search/test_elasticsearch.py +++ b/ingestion/tests/unit/topology/search/test_elasticsearch.py @@ -194,7 +194,7 @@ class ElasticSearchUnitTest(TestCase): def __init__(self, methodName, test_connection) -> None: super().__init__(methodName) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(mock_es_config) + self.config = OpenMetadataWorkflowConfig.model_validate(mock_es_config) self.es_source = ElasticsearchSource.create( mock_es_config["source"], self.config.workflowConfig.openMetadataServerConfig, diff --git a/ingestion/tests/unit/topology/storage/test_storage.py b/ingestion/tests/unit/topology/storage/test_storage.py index 0497fd8fa35e..9c3526e7ef3a 100644 --- a/ingestion/tests/unit/topology/storage/test_storage.py +++ b/ingestion/tests/unit/topology/storage/test_storage.py @@ -143,7 +143,9 @@ class StorageUnitTest(TestCase): def __init__(self, method_name: str, test_connection) -> None: super().__init__(method_name) test_connection.return_value = False - self.config = OpenMetadataWorkflowConfig.parse_obj(MOCK_OBJECT_STORE_CONFIG) + self.config = OpenMetadataWorkflowConfig.model_validate( + MOCK_OBJECT_STORE_CONFIG + ) # This already validates that the source can be initialized self.object_store_source = S3Source.create( @@ -430,5 +432,5 @@ def return_metadata_entry(self): verbose=False, ) content = json.loads(metadata_config_response) - container_config = StorageContainerConfig.parse_obj(content) + container_config = StorageContainerConfig.model_validate(content) return container_config.entries diff --git a/openmetadata-airflow-apis/openmetadata_managed_apis/operations/deploy.py b/openmetadata-airflow-apis/openmetadata_managed_apis/operations/deploy.py index b35141f6f31f..5996bfa98226 100644 --- a/openmetadata-airflow-apis/openmetadata_managed_apis/operations/deploy.py +++ b/openmetadata-airflow-apis/openmetadata_managed_apis/operations/deploy.py @@ -35,7 +35,6 @@ from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import ( IngestionPipeline, ) -from metadata.ingestion.models.encoders import show_secrets_encoder from metadata.utils.secrets.secrets_manager_factory import SecretsManagerFactory logger = operations_logger() @@ -75,7 +74,7 @@ def store_airflow_pipeline_config( logger.info(f"Saving file to {dag_config_file_path}") with open(dag_config_file_path, "w") as outfile: - outfile.write(self.ingestion_pipeline.json(encoder=show_secrets_encoder)) + outfile.write(self.ingestion_pipeline.model_dump_json()) return {"workflow_config_file": str(dag_config_file_path)} diff --git a/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/application.py b/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/application.py index 00e5f9f00a57..99678759f1f5 100644 --- a/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/application.py +++ b/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/application.py @@ -33,7 +33,6 @@ from metadata.generated.schema.metadataIngestion.applicationPipeline import ( ApplicationPipeline, ) -from metadata.ingestion.models.encoders import show_secrets_encoder from metadata.workflow.application import ApplicationWorkflow from metadata.workflow.application_output_handler import print_status @@ -50,7 +49,7 @@ def application_workflow(workflow_config: OpenMetadataApplicationConfig): set_operator_logger(workflow_config) - config = json.loads(workflow_config.json(encoder=show_secrets_encoder)) + config = json.loads(workflow_config.model_dump_json()) workflow = ApplicationWorkflow.create(config) workflow.execute() diff --git a/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/common.py b/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/common.py index 3f3d7172990b..e2f2150b13e6 100644 --- a/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/common.py +++ b/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/common.py @@ -17,7 +17,6 @@ from functools import partial from typing import Callable, Optional, Union -import airflow from airflow import DAG from openmetadata_managed_apis.api.utils import clean_dag_id from pydantic import ValidationError @@ -34,7 +33,7 @@ from metadata.generated.schema.metadataIngestion.application import ( OpenMetadataApplicationConfig, ) -from metadata.ingestion.models.encoders import show_secrets_encoder +from metadata.generated.schema.type.basic import Timestamp, Uuid from metadata.ingestion.ometa.ometa_api import OpenMetadata from metadata.utils import fqn from metadata.workflow.workflow_output_handler import print_status @@ -203,7 +202,7 @@ def metadata_ingestion_workflow(workflow_config: OpenMetadataWorkflowConfig): set_operator_logger(workflow_config) - config = json.loads(workflow_config.json(encoder=show_secrets_encoder)) + config = json.loads(workflow_config.model_dump_json()) workflow = MetadataWorkflow.create(config) workflow.execute() @@ -247,14 +246,18 @@ def build_dag_configs(ingestion_pipeline: IngestionPipeline) -> dict: :param ingestion_pipeline: pipeline configs :return: dict to use as kwargs """ + + if ingestion_pipeline.airflowConfig.startDate: + start_date = ingestion_pipeline.airflowConfig.startDate.root + else: + start_date = datetime.now() - timedelta(days=1) + return { "dag_id": clean_dag_id(ingestion_pipeline.name.root), "description": ingestion_pipeline.description.root if ingestion_pipeline.description is not None else None, - "start_date": ingestion_pipeline.airflowConfig.startDate.root - if ingestion_pipeline.airflowConfig.startDate - else airflow.utils.dates.days_ago(1), + "start_date": start_date, "end_date": ingestion_pipeline.airflowConfig.endDate.root if ingestion_pipeline.airflowConfig.endDate else None, @@ -312,7 +315,7 @@ def send_failed_status_callback(workflow_config: OpenMetadataWorkflowConfig, *_, pipeline_status = metadata.get_pipeline_status( workflow_config.ingestionPipelineFQN, str(workflow_config.pipelineRunId) ) - pipeline_status.endDate = datetime.now().timestamp() * 1000 + pipeline_status.endDate = Timestamp(int(datetime.now().timestamp() * 1000)) pipeline_status.pipelineState = PipelineState.failed metadata.create_or_update_pipeline_status( @@ -350,7 +353,7 @@ def build_dag( with DAG(**build_dag_configs(ingestion_pipeline)) as dag: # Initialize with random UUID4. Will be used by the callback instead of # generating it inside the Workflow itself. - workflow_config.pipelineRunId = str(uuid.uuid4()) + workflow_config.pipelineRunId = Uuid(uuid.uuid4()) CustomPythonOperator( task_id=task_name, diff --git a/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/data_insight.py b/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/data_insight.py index 2d5d633ff2f8..820654160a03 100644 --- a/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/data_insight.py +++ b/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/data_insight.py @@ -40,7 +40,6 @@ SourceConfig, WorkflowConfig, ) -from metadata.ingestion.models.encoders import show_secrets_encoder from metadata.ingestion.ometa.ometa_api import OpenMetadata from metadata.workflow.data_insight import DataInsightWorkflow from metadata.workflow.workflow_output_handler import print_status @@ -59,7 +58,7 @@ def data_insight_workflow(workflow_config: OpenMetadataWorkflowConfig): """ set_operator_logger(workflow_config) - config = json.loads(workflow_config.json(encoder=show_secrets_encoder)) + config = json.loads(workflow_config.model_dump_json()) workflow = DataInsightWorkflow.create(config) workflow.execute() diff --git a/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/profiler.py b/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/profiler.py index 5810c363b762..62822387961c 100644 --- a/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/profiler.py +++ b/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/profiler.py @@ -27,7 +27,6 @@ Sink, WorkflowConfig, ) -from metadata.ingestion.models.encoders import show_secrets_encoder from metadata.workflow.profiler import ProfilerWorkflow from metadata.workflow.workflow_output_handler import print_status @@ -44,7 +43,7 @@ def profiler_workflow(workflow_config: OpenMetadataWorkflowConfig): set_operator_logger(workflow_config) - config = json.loads(workflow_config.json(encoder=show_secrets_encoder)) + config = json.loads(workflow_config.model_dump_json()) workflow = ProfilerWorkflow.create(config) workflow.execute() diff --git a/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/test_suite.py b/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/test_suite.py index bf282263aec6..37d5d4314e2a 100644 --- a/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/test_suite.py +++ b/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/test_suite.py @@ -27,7 +27,6 @@ Sink, WorkflowConfig, ) -from metadata.ingestion.models.encoders import show_secrets_encoder from metadata.workflow.data_quality import TestSuiteWorkflow from metadata.workflow.workflow_output_handler import print_status @@ -44,7 +43,7 @@ def test_suite_workflow(workflow_config: OpenMetadataWorkflowConfig): set_operator_logger(workflow_config) - config = json.loads(workflow_config.json(encoder=show_secrets_encoder)) + config = json.loads(workflow_config.model_dump_json()) workflow = TestSuiteWorkflow.create(config) workflow.execute() diff --git a/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/usage.py b/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/usage.py index 2dd3a3772641..77f71281c4ed 100644 --- a/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/usage.py +++ b/openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/usage.py @@ -31,7 +31,6 @@ Processor, Stage, ) -from metadata.ingestion.models.encoders import show_secrets_encoder from metadata.workflow.usage import UsageWorkflow from metadata.workflow.workflow_output_handler import print_status @@ -48,7 +47,7 @@ def usage_workflow(workflow_config: OpenMetadataWorkflowConfig): set_operator_logger(workflow_config) - config = json.loads(workflow_config.json(encoder=show_secrets_encoder)) + config = json.loads(workflow_config.model_dump_json()) workflow = UsageWorkflow.create(config) workflow.execute() diff --git a/openmetadata-airflow-apis/tests/unit/ingestion_pipeline/test_workflow_creation.py b/openmetadata-airflow-apis/tests/unit/ingestion_pipeline/test_workflow_creation.py index 0ab1a4ff7474..a056e13a77e0 100644 --- a/openmetadata-airflow-apis/tests/unit/ingestion_pipeline/test_workflow_creation.py +++ b/openmetadata-airflow-apis/tests/unit/ingestion_pipeline/test_workflow_creation.py @@ -14,6 +14,7 @@ """ import json import uuid +from datetime import datetime from unittest import TestCase from unittest.mock import patch @@ -64,9 +65,9 @@ from metadata.generated.schema.security.client.openMetadataJWTClientConfig import ( OpenMetadataJWTClientConfig, ) +from metadata.generated.schema.type.basic import DateTime from metadata.generated.schema.type.entityReference import EntityReference from metadata.ingestion.api.parser import parse_workflow_config_gracefully -from metadata.ingestion.models.encoders import show_secrets_encoder from metadata.ingestion.ometa.ometa_api import OpenMetadata from metadata.workflow.data_quality import TestSuiteWorkflow from metadata.workflow.metadata import MetadataWorkflow @@ -77,6 +78,9 @@ def mock_set_ingestion_pipeline_status(self, state): return True +START_DATE = DateTime(datetime.strptime("2022-06-10T15:06:47", "%Y-%m-%dT%H:%M:%S")) + + class OMetaServiceTest(TestCase): """ Run this integration test with the local API available @@ -190,7 +194,7 @@ def test_ingestion_workflow(self): sourceConfig=SourceConfig(config=DatabaseServiceMetadataPipeline()), openMetadataServerConnection=self.server_config, airflowConfig=AirflowConfig( - startDate="2022-06-10T15:06:47+00:00", + startDate=START_DATE, ), service=EntityReference( id=self.service.id, @@ -200,7 +204,7 @@ def test_ingestion_workflow(self): ) workflow_config = build_metadata_workflow_config(ingestion_pipeline) - config = json.loads(workflow_config.json(encoder=show_secrets_encoder)) + config = json.loads(workflow_config.model_dump_json()) parse_workflow_config_gracefully(config) @@ -223,7 +227,7 @@ def test_usage_workflow(self): sourceConfig=SourceConfig(config=DatabaseServiceQueryUsagePipeline()), openMetadataServerConnection=self.server_config, airflowConfig=AirflowConfig( - startDate="2022-06-10T15:06:47+00:00", + startDate=START_DATE, ), service=EntityReference( id=self.usage_service.id, @@ -235,7 +239,7 @@ def test_usage_workflow(self): workflow_config = build_usage_workflow_config(ingestion_pipeline) self.assertIn("usage", workflow_config.source.type) - config = json.loads(workflow_config.json(encoder=show_secrets_encoder)) + config = json.loads(workflow_config.model_dump_json()) parse_workflow_config_gracefully(config) @@ -258,7 +262,7 @@ def test_lineage_workflow(self): sourceConfig=SourceConfig(config=DatabaseServiceQueryLineagePipeline()), openMetadataServerConnection=self.server_config, airflowConfig=AirflowConfig( - startDate="2022-06-10T15:06:47+00:00", + startDate=START_DATE, ), service=EntityReference( id=self.usage_service.id, @@ -270,7 +274,7 @@ def test_lineage_workflow(self): workflow_config = build_lineage_workflow_config(ingestion_pipeline) self.assertIn("lineage", workflow_config.source.type) - config = json.loads(workflow_config.json(encoder=show_secrets_encoder)) + config = json.loads(workflow_config.model_dump_json()) parse_workflow_config_gracefully(config) @@ -293,7 +297,7 @@ def test_profiler_workflow(self): sourceConfig=SourceConfig(config=DatabaseServiceProfilerPipeline()), openMetadataServerConnection=self.server_config, airflowConfig=AirflowConfig( - startDate="2022-06-10T15:06:47+00:00", + startDate=START_DATE, ), service=EntityReference( id=self.service.id, @@ -303,7 +307,7 @@ def test_profiler_workflow(self): ) workflow_config = build_profiler_workflow_config(ingestion_pipeline) - config = json.loads(workflow_config.json(encoder=show_secrets_encoder)) + config = json.loads(workflow_config.model_dump_json()) parse_workflow_config_gracefully(config) @@ -331,7 +335,7 @@ def test_test_suite_workflow(self): ), openMetadataServerConnection=self.server_config, airflowConfig=AirflowConfig( - startDate="2022-06-10T15:06:47+00:00", + startDate=START_DATE, ), service=EntityReference( id=uuid.uuid4(), @@ -341,6 +345,6 @@ def test_test_suite_workflow(self): ) workflow_config = build_test_suite_workflow_config(ingestion_pipeline) - config = json.loads(workflow_config.json(encoder=show_secrets_encoder)) + config = json.loads(workflow_config.model_dump_json()) parse_workflow_config_gracefully(config) diff --git a/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/metadata/openMetadataConnection.json b/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/metadata/openMetadataConnection.json index fb6c56cfd29b..24bb3ce30e2c 100644 --- a/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/metadata/openMetadataConnection.json +++ b/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/metadata/openMetadataConnection.json @@ -160,7 +160,7 @@ "limitRecords": { "description": "Limit the number of records for Indexing.", "type": "integer", - "default": "1000" + "default": 1000 }, "forceEntityOverwriting": { "description": "Force the overwriting of any entity during the ingestion.", diff --git a/scripts/datamodel_generation.py b/scripts/datamodel_generation.py index d30537ba5f23..085d3231d16f 100644 --- a/scripts/datamodel_generation.py +++ b/scripts/datamodel_generation.py @@ -81,3 +81,21 @@ with open(file_path, "w", encoding=UTF_8) as file_: file_.write(content) +# Until https://github.com/koxudaxi/datamodel-code-generator/issues/1996 +# Supporting timezone aware datetime is too complex for the profiler +DATETIME_AWARE_FILE_PATHS = [ + f"{ingestion_path}src/metadata/generated/schema/type/basic.py", + # f"{ingestion_path}src/metadata/generated/schema/entity/data/table.py", +] + +for file_path in DATETIME_AWARE_FILE_PATHS: + with open(file_path, "r", encoding=UTF_8) as file_: + content = file_.read() + content = content.replace( + "from pydantic import AnyUrl, AwareDatetime, ConfigDict, EmailStr, Field, RootModel", + "from pydantic import AnyUrl, ConfigDict, EmailStr, Field, RootModel" + ) + content = content.replace("from datetime import date, time", "from datetime import date, time, datetime") + content = content.replace("AwareDatetime", "datetime") + with open(file_path, "w", encoding=UTF_8) as file_: + file_.write(content)