yujunjun/OpenMetadata
1
0
Fork 0
mirror of https://github.com/open-metadata/OpenMetadata.git synced 2025-12-05 03:54:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

#15243 - Pydantic V2 & Airflow 2.9 (#16480)

Browse Source 
* pydantic v2

* pydanticv2

* fix parser

* fix annotated

* fix model dumping

* mysql ingestion

* clean root models

* clean root models

* bump airflow

* bump airflow

* bump airflow

* optionals

* optionals

* optionals

* jdk

* airflow migrate

* fab provider

* fab provider

* fab provider

* some more fixes

* fixing tests and imports

* model_dump and model_validate

* model_dump and model_validate

* model_dump and model_validate

* union

* pylint

* pylint

* integration tests

* fix CostAnalysisReportData

* integration tests

* tests

* missing defaults

* missing defaults
This commit is contained in:
Pere Miquel Brull 2024-06-05 21:18:37 +02:00 committed by GitHub
parent fcb87b5866
commit d8e2187980
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
521 changed files with 4326 additions and 3701 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
ingestion/Dockerfile
View File

@ -1,6 +1,6 @@
FROM mysql:8.3 as mysql
FROM apache/airflow:2.7.3-python3.10
FROM apache/airflow:2.9.1-python3.10
USER root
RUN curl -sS https://packages.microsoft.com/keys/microsoft.asc | apt-key add -
RUN curl -sS https://packages.microsoft.com/config/debian/11/prod.list > /etc/apt/sources.list.d/mssql-release.list
@ -27,13 +27,13 @@ RUN apt-get -qq update \
libssl-dev \
libxml2 \
libkrb5-dev \
openjdk-11-jre \
default-jdk \
openssl \
postgresql \
postgresql-contrib \
tdsodbc \
unixodbc \
unixodbc-dev \
unixodbc=2.3.11-2+deb12u1 \
unixodbc-dev=2.3.11-2+deb12u1 \
unzip \
git \
wget --no-install-recommends \
@ -53,22 +53,6 @@ RUN if [[ $(uname -m) == "arm64" || $(uname -m) == "aarch64" ]]; \
ENV LD_LIBRARY_PATH=/instantclient
# Security patches for base image
# monitor no fixed version for
# https://security.snyk.io/vuln/SNYK-DEBIAN11-LIBTASN16-3061097
# https://security.snyk.io/vuln/SNYK-DEBIAN11-MARIADB105-2940589
# https://security.snyk.io/vuln/SNYK-DEBIAN11-BIND9-3027852
# https://security.snyk.io/vuln/SNYK-DEBIAN11-EXPAT-3023031 we are already installed the latest
RUN echo "deb http://deb.debian.org/debian bullseye-backports main" > /etc/apt/sources.list.d/backports.list
RUN apt-get -qq update \
&& apt-get -qq install -t bullseye-backports -y \
curl \
libpcre2-8-0 \
postgresql-common \
expat \
bind9 \
&& rm -rf /var/lib/apt/lists/*
# Required for Starting Ingestion Container in Docker Compose
COPY --chown=airflow:0 --chmod=775 ingestion/ingestion_dependency.sh /opt/airflow
# Required for Ingesting Sample Data
@ -99,7 +83,7 @@ RUN if [[ $(uname -m) != "aarch64" ]]; \
# bump python-daemon for https://github.com/apache/airflow/pull/29916
RUN pip install "python-daemon>=3.0.0"
# remove all airflow providers except for docker and cncf kubernetes
RUN pip freeze | grep "apache-airflow-providers" | grep --invert-match -E "docker|http|cncf" | xargs pip uninstall -y
RUN pip freeze | grep "apache-airflow-providers" | grep --invert-match -E "docker|http|cncf|fab" | xargs pip uninstall -y
# Uninstalling psycopg2-binary and installing psycopg2 instead
# because the psycopg2-binary generates a architecture specific error
# while authenticating connection with the airflow, psycopg2 solves this error
@ -108,4 +92,4 @@ RUN pip install psycopg2 mysqlclient==2.1.1
# Make required folders for openmetadata-airflow-apis
RUN mkdir -p /opt/airflow/dag_generated_configs
# This is required as it's responsible to create airflow.cfg file
RUN airflow db init && rm -f /opt/airflow/airflow.db
RUN airflow db migrate && rm -f /opt/airflow/airflow.db

30
ingestion/Dockerfile.ci
View File

@ -1,11 +1,12 @@
FROM mysql:8.3 as mysql
FROM apache/airflow:2.7.3-python3.10
FROM apache/airflow:2.9.1-python3.10
USER root
RUN curl -sS https://packages.microsoft.com/keys/microsoft.asc | apt-key add -
RUN curl -sS https://packages.microsoft.com/config/debian/11/prod.list > /etc/apt/sources.list.d/mssql-release.list
# Install Dependencies (listed in alphabetical order)
RUN apt-get -qq update \
RUN dpkg --configure -a \
&& apt-get -qq update \
&& apt-get -qq install -y \
alien \
build-essential \
@ -26,13 +27,13 @@ RUN apt-get -qq update \
libssl-dev \
libxml2 \
libkrb5-dev \
openjdk-11-jre \
default-jdk \
openssl \
postgresql \
postgresql-contrib \
tdsodbc \
unixodbc \
unixodbc-dev \
unixodbc=2.3.11-2+deb12u1 \
unixodbc-dev=2.3.11-2+deb12u1 \
unzip \
vim \
git \
@ -53,21 +54,6 @@ RUN if [[ $(uname -m) == "arm64" || $(uname -m) == "aarch64" ]]; \
ENV LD_LIBRARY_PATH=/instantclient
# Security patches for base image
# monitor no fixed version for
# https://security.snyk.io/vuln/SNYK-DEBIAN11-LIBTASN16-3061097
# https://security.snyk.io/vuln/SNYK-DEBIAN11-MARIADB105-2940589
# https://security.snyk.io/vuln/SNYK-DEBIAN11-BIND9-3027852
# https://security.snyk.io/vuln/SNYK-DEBIAN11-EXPAT-3023031 we are already installed the latest
RUN echo "deb http://deb.debian.org/debian bullseye-backports main" > /etc/apt/sources.list.d/backports.list
RUN apt-get -qq update \
&& apt-get -qq install -t bullseye-backports -y \
curl \
libpcre2-8-0 \
postgresql-common \
expat \
bind9
# Required for Starting Ingestion Container in Docker Compose
# Provide Execute Permissions to shell script
COPY --chown=airflow:0 --chmod=775 ingestion/ingestion_dependency.sh /opt/airflow
@ -109,7 +95,7 @@ RUN if [[ $(uname -m) != "aarch64" ]]; \
RUN pip install "python-daemon>=3.0.0"
# remove all airflow providers except for docker and cncf kubernetes
RUN pip freeze | grep "apache-airflow-providers" | grep --invert-match -E "docker|http|cncf" | xargs pip uninstall -y
RUN pip freeze | grep "apache-airflow-providers" | grep --invert-match -E "docker|http|cncf|fab" | xargs pip uninstall -y
# Uninstalling psycopg2-binary and installing psycopg2 instead
# because the psycopg2-binary generates a architecture specific error
@ -121,4 +107,4 @@ RUN mkdir -p /opt/airflow/dag_generated_configs
EXPOSE 8080
# This is required as it's responsible to create airflow.cfg file
RUN airflow db init && rm -f /opt/airflow/airflow.db
RUN airflow db migrate && rm -f /opt/airflow/airflow.db

2
ingestion/examples/airflow/dags/airflow_docker_operator.py
View File

@ -61,7 +61,7 @@ with models.DAG(
environment={"config": config, "pipelineType": PipelineType.metadata.value},
docker_url="unix://var/run/docker.sock", # To allow to start Docker. Needs chmod 666 permissions
tty=True,
auto_remove="True",
auto_remove="success",
network_mode="host", # To reach the OM server
task_id="ingest",
dag=dag,

30
ingestion/examples/sample_data/datasets/tableTests.json
View File

@ -6,7 +6,7 @@
"description": "Rows should always be 100 because of something",
"testCase": {
"config": {
"value": 120
"value": "120"
},
"tableTestType": "tableRowCountToEqual"
},
@ -21,7 +21,7 @@
"description": "Rows should always be 100 because of something",
"testCase": {
"config": {
"value": 120
"value": "120"
},
"tableTestType": "tableRowCountToEqual"
},
@ -36,7 +36,7 @@
"description": "We expect certain columns",
"testCase": {
"config": {
"value": 5
"value": "5"
},
"tableTestType": "tableColumnCountToEqual"
},
@ -51,8 +51,8 @@
"description": "Rows should always be 100 because of something",
"testCase": {
"config": {
"minValue": 100,
"maxValue": 200
"minValue": "100",
"maxValue": "200"
},
"tableTestType": "tableRowCountToBeBetween"
},
@ -67,8 +67,8 @@
"description": "Rows should always be 100 because of something",
"testCase": {
"config": {
"minValue": 100,
"maxValue": 200
"minValue": "100",
"maxValue": "200"
},
"tableTestType": "tableRowCountToBeBetween"
},
@ -86,7 +86,7 @@
"description": "user_id should be positive",
"testCase": {
"config": {
"minValue": 0
"minValue": "0"
},
"columnTestType": "columnValuesToBeBetween"
},
@ -102,7 +102,7 @@
"description": "user_id should be positive",
"testCase": {
"config": {
"minValue": 0
"minValue": "0"
},
"columnTestType": "columnValuesToBeBetween"
},
@ -206,7 +206,7 @@
"description": "Some description...",
"testCase": {
"config": {
"missingCountValue": 10
"missingCountValue": "10"
},
"columnTestType": "columnValuesMissingCountToBeEqual"
},
@ -222,7 +222,7 @@
"description": "Some description...",
"testCase": {
"config": {
"missingCountValue": 10
"missingCountValue": "10"
},
"columnTestType": "columnValuesMissingCountToBeEqual"
},
@ -238,8 +238,8 @@
"description": "email should have a fixed length",
"testCase": {
"config": {
"minValue": 6,
"maxValue": 30
"minValue": "6",
"maxValue": "30"
},
"columnTestType": "columnValuesToBeBetween"
},
@ -255,8 +255,8 @@
"description": "email should have a fixed length",
"testCase": {
"config": {
"minValue": 6,
"maxValue": 30
"minValue": "6",
"maxValue": "30"
},
"columnTestType": "columnValuesToBeBetween"
},

18
ingestion/examples/sample_data/tests/testSuites.json
View File

@ -14,7 +14,7 @@
"parameterValues": [
{
"name": "columnCount",
"value": 10
"value": "10"
}
],
"resolutions": {
@ -95,11 +95,11 @@
"parameterValues": [
{
"name": "minColValue",
"value": 1
"value": "1"
},
{
"name": "maxColValue",
"value": 10
"value": "10"
}
],
"resolutions": {
@ -169,11 +169,11 @@
"parameterValues": [
{
"name": "minValueForMaxInCol",
"value": 50
"value": "50"
},
{
"name": "maxValueForMaxInCol",
"value": 100
"value": "100"
}
],
"resolutions": {
@ -243,11 +243,11 @@
"parameterValues": [
{
"name": "min",
"value": 90001
"value": "90001"
},
{
"name": "max",
"value": 96162
"value": "96162"
}
],
"resolutions": {
@ -332,11 +332,11 @@
"parameterValues": [
{
"name": "min",
"value": 90001
"value": "90001"
},
{
"name": "max",
"value": 96162
"value": "96162"
}
],
"resolutions": {}

5
ingestion/ingestion_dependency.sh
View File

@ -33,7 +33,7 @@ export AIRFLOW__API__AUTH_BACKEND=${AIRFLOW__API__AUTH_BACKENDS:-"airflow.api.au
# Use the default airflow env var or the one we set from OM properties
export AIRFLOW__DATABASE__SQL_ALCHEMY_CONN=${AIRFLOW__DATABASE__SQL_ALCHEMY_CONN:-$DB_CONN}
airflow db init
airflow db migrate
airflow users create \
--username ${AIRFLOW_ADMIN_USER} \
@ -43,9 +43,6 @@ airflow users create \
--email spiderman@superhero.org \
--password ${AIRFLOW_ADMIN_PASSWORD}
(sleep 5; airflow db migrate)
(sleep 5; airflow db migrate)
# we need to this in case the container is restarted and the scheduler exited without tidying up its lock file
rm -f /opt/airflow/airflow-webserver-monitor.pid
airflow webserver --port 8080 -D &

20
ingestion/operators/docker/Dockerfile
View File

@ -4,7 +4,8 @@ RUN curl -sS https://packages.microsoft.com/keys/microsoft.asc | apt-key add -
RUN curl -sS https://packages.microsoft.com/config/debian/11/prod.list > /etc/apt/sources.list.d/mssql-release.list
# Install Dependencies (listed in alphabetical order)
RUN apt-get -qq update \
RUN dpkg --configure -a \
&& apt-get -qq update \
&& apt-get -qq install -y \
alien \
build-essential \
@ -25,7 +26,7 @@ RUN apt-get -qq update \
libssl-dev \
libxml2 \
libkrb5-dev \
openjdk-11-jre \
default-jdk \
openssl \
postgresql \
postgresql-contrib \
@ -58,21 +59,6 @@ RUN if [[ $(uname -m) == "arm64" || $(uname -m) == "aarch64" ]]; \
ENV LD_LIBRARY_PATH=/instantclient
# Security patches for base image
# monitor no fixed version for
# https://security.snyk.io/vuln/SNYK-DEBIAN11-LIBTASN16-3061097
# https://security.snyk.io/vuln/SNYK-DEBIAN11-MARIADB105-2940589
# https://security.snyk.io/vuln/SNYK-DEBIAN11-BIND9-3027852
# https://security.snyk.io/vuln/SNYK-DEBIAN11-EXPAT-3023031 we are already installed the latest
RUN echo "deb http://deb.debian.org/debian bullseye-backports main" > /etc/apt/sources.list.d/backports.list
RUN apt-get -qq update \
&& apt-get install -t bullseye-backports -y \
curl \
libpcre2-8-0 \
postgresql-common \
expat \
bind9
WORKDIR ingestion/
# Required for Airflow DockerOperator, as we need to run the workflows from a `python main.py` command in the container.

17
ingestion/operators/docker/Dockerfile.ci
View File

@ -25,7 +25,7 @@ RUN apt-get -qq update \
libssl-dev \
libxml2 \
libkrb5-dev \
openjdk-11-jre \
default-jdk \
openssl \
postgresql \
postgresql-contrib \
@ -59,21 +59,6 @@ RUN if [[ $(uname -m) == "arm64" || $(uname -m) == "aarch64" ]]; \
ENV LD_LIBRARY_PATH=/instantclient
# Security patches for base image
# monitor no fixed version for
# https://security.snyk.io/vuln/SNYK-DEBIAN11-LIBTASN16-3061097
# https://security.snyk.io/vuln/SNYK-DEBIAN11-MARIADB105-2940589
# https://security.snyk.io/vuln/SNYK-DEBIAN11-BIND9-3027852
# https://security.snyk.io/vuln/SNYK-DEBIAN11-EXPAT-3023031 we are already installed the latest
RUN echo "deb http://deb.debian.org/debian bullseye-backports main" > /etc/apt/sources.list.d/backports.list
RUN apt-get -qq update \
&& apt-get -qq install -t bullseye-backports -y \
curl \
libpcre2-8-0 \
postgresql-common \
expat \
bind9
WORKDIR ingestion/
# For the dev build, we copy all files

4
ingestion/operators/docker/exit_handler.py
View File

@ -86,13 +86,13 @@ def main():
pipeline_status = metadata.get_pipeline_status(
workflow_config.ingestionPipelineFQN,
str(workflow_config.pipelineRunId.__root__),
str(workflow_config.pipelineRunId.root),
)
# Maybe the workflow was not even initialized
if not pipeline_status:
pipeline_status = PipelineStatus(
runId=str(workflow_config.pipelineRunId.__root__),
runId=str(workflow_config.pipelineRunId.root),
startDate=datetime.now().timestamp() * 1000,
timestamp=datetime.now().timestamp() * 1000,
)

45
ingestion/setup.py
View File

@ -19,31 +19,31 @@ from setuptools import setup
# Add here versions required for multiple plugins
VERSIONS = {
"airflow": "apache-airflow==2.7.3",
"adlfs": "adlfs~=2022.11",
"airflow": "apache-airflow==2.9.1",
"adlfs": "adlfs>=2023.1.0",
"avro": "avro>=1.11.3,<1.12",
"boto3": "boto3>=1.20,<2.0", # No need to add botocore separately. It's a dep from boto3
"geoalchemy2": "GeoAlchemy2~=0.12",
"google-cloud-storage": "google-cloud-storage==1.43.0",
"gcsfs": "gcsfs~=2022.11",
"gcsfs": "gcsfs>=2023.1.0",
"great-expectations": "great-expectations>=0.18.0,<0.18.14",
"grpc-tools": "grpcio-tools>=1.47.2",
"msal": "msal~=1.2",
"neo4j": "neo4j~=5.3.0",
"pandas": "pandas~=2.0.0",
"pyarrow": "pyarrow~=14.0",
"pydantic": "pydantic~=1.10",
"pyarrow": "pyarrow~=16.0",
"pydantic": "pydantic~=2.0",
"pydomo": "pydomo~=0.3",
"pymysql": "pymysql>=1.0.2",
"pymysql": "pymysql~=1.0",
"pyodbc": "pyodbc>=4.0.35,<5",
"scikit-learn": "scikit-learn~=1.0", # Python 3.7 only goes up to 1.0.2
"packaging": "packaging==21.3",
"packaging": "packaging",
"azure-storage-blob": "azure-storage-blob~=12.14",
"azure-identity": "azure-identity~=1.12",
"sqlalchemy-databricks": "sqlalchemy-databricks~=0.1",
"databricks-sdk": "databricks-sdk>=0.18.0,<0.20.0",
"trino": "trino[sqlalchemy]",
"spacy": "spacy==3.5.0",
"spacy": "spacy~=3.7",
"looker-sdk": "looker-sdk>=22.20.0",
"lkml": "lkml~=1.3",
"tableau": "tableau-api-lib~=0.1",
@ -99,7 +99,7 @@ base_requirements = {
"cached-property==1.5.2", # LineageParser
"chardet==4.0.0", # Used in the profiler
"cryptography>=42.0.0",
"email-validator>=1.0.3", # For the pydantic generated models for Email
"email-validator>=2.0", # For the pydantic generated models for Email
"importlib-metadata>=4.13.0", # From airflow constraints
"Jinja2>=2.11.3",
"jsonpatch<2.0, >=1.24",
@ -125,7 +125,7 @@ plugins: Dict[str, Set[str]] = {
"attrs",
}, # Same as ingestion container. For development.
"amundsen": {VERSIONS["neo4j"]},
"athena": {"pyathena==3.0.8"},
"athena": {"pyathena~=3.0"},
"atlas": {},
"azuresql": {VERSIONS["pyodbc"]},
"azure-sso": {VERSIONS["msal"]},
@ -165,7 +165,7 @@ plugins: Dict[str, Set[str]] = {
"datalake-azure": {
VERSIONS["azure-storage-blob"],
VERSIONS["azure-identity"],
VERSIONS["adlfs"], # Python 3.7 does only support up to 2022.2.0
VERSIONS["adlfs"],
*COMMONS["datalake"],
},
"datalake-gcs": {
@ -178,7 +178,7 @@ plugins: Dict[str, Set[str]] = {
# https://github.com/fsspec/s3fs/blob/9bf99f763edaf7026318e150c4bd3a8d18bb3a00/requirements.txt#L1
# however, the latest version of `s3fs` conflicts its `aiobotocore` dep with `boto3`'s dep on `botocore`.
# Leaving this marked to the automatic resolution to speed up installation.
"s3fs==0.4.2",
"s3fs",
*COMMONS["datalake"],
},
"deltalake": {"delta-spark<=2.3.0"},
@ -201,7 +201,7 @@ plugins: Dict[str, Set[str]] = {
"impyla~=0.18.0",
},
"iceberg": {
"pyiceberg<1",
"pyiceberg>=0.5",
# Forcing the version of a few packages so it plays nicely with other requirements.
VERSIONS["pydantic"],
VERSIONS["adlfs"],
@ -224,7 +224,7 @@ plugins: Dict[str, Set[str]] = {
"gitpython~=3.1.34",
VERSIONS["giturlparse"],
},
"mlflow": {"mlflow-skinny>=2.3.0", "alembic~=1.10.2"},
"mlflow": {"mlflow-skinny>=2.3.0"},
"mongo": {VERSIONS["mongo"], VERSIONS["pandas"]},
"couchbase": {"couchbase~=4.1"},
"mssql": {"sqlalchemy-pytds~=0.3"},
@ -234,7 +234,7 @@ plugins: Dict[str, Set[str]] = {
"openlineage": {*COMMONS["kafka"]},
"oracle": {"cx_Oracle>=8.3.0,<9", "oracledb~=1.2"},
"pgspider": {"psycopg2-binary", "sqlalchemy-pgspider"},
"pinotdb": {"pinotdb~=0.3"},
"pinotdb": {"pinotdb~=5.0"},
"postgres": {*COMMONS["postgres"]},
"powerbi": {
VERSIONS["msal"],
@ -256,7 +256,7 @@ plugins: Dict[str, Set[str]] = {
VERSIONS["geoalchemy2"],
},
"sagemaker": {VERSIONS["boto3"]},
"salesforce": {"simple_salesforce==1.11.4"},
"salesforce": {"simple_salesforce~=1.11"},
"sample-data": {VERSIONS["avro"], VERSIONS["grpc-tools"]},
"sap-hana": {"hdbcli", "sqlalchemy-hana"},
"sas": {},
@ -277,12 +277,13 @@ plugins: Dict[str, Set[str]] = {
dev = {
"black==22.3.0",
"datamodel-code-generator==0.24.2",
"boto3-stubs[essential]",
"datamodel-code-generator==0.25.6",
"boto3-stubs",
"mypy-boto3-glue",
"isort",
"pre-commit",
"pycln",
"pylint~=3.0.0",
"pylint~=3.0",
# For publishing
"twine",
"build",
@ -293,11 +294,12 @@ dev = {
test = {
# Install Airflow as it's not part of `all` plugin
VERSIONS["airflow"],
"boto3-stubs[boto3]",
"boto3-stubs",
"mypy-boto3-glue",
"coverage",
# Install GE because it's not in the `all` plugin
VERSIONS["great-expectations"],
"moto==4.0.8",
"moto~=5.0",
"pytest==7.0.0",
"pytest-cov",
"pytest-order",
@ -326,6 +328,7 @@ test = {
"minio==7.2.5",
*plugins["mlflow"],
*plugins["datalake-s3"],
*plugins["pii-processor"],
"requests==2.31.0",
}

6
ingestion/src/airflow_provider_openmetadata/lineage/callback.py
View File

@ -48,7 +48,7 @@ def failure_callback(context: Dict[str, str]) -> None:
)
pipeline: Pipeline = metadata.get_by_name(
entity=Pipeline,
fqn=f"{airflow_service_entity.name.__root__}.{dag.dag_id}",
fqn=f"{airflow_service_entity.name.root}.{dag.dag_id}",
)
if pipeline:
@ -60,7 +60,7 @@ def failure_callback(context: Dict[str, str]) -> None:
)
else:
logging.warning(
f"Pipeline {airflow_service_entity.name.__root__}.{dag.dag_id} not found. Skipping status update."
f"Pipeline {airflow_service_entity.name.root}.{dag.dag_id} not found. Skipping status update."
)
except Exception as exc: # pylint: disable=broad-except
@ -90,7 +90,7 @@ def success_callback(context: Dict[str, str]) -> None:
)
pipeline: Pipeline = metadata.get_by_name(
entity=Pipeline,
fqn=f"{airflow_service_entity.name.__root__}.{dag.dag_id}",
fqn=f"{airflow_service_entity.name.root}.{dag.dag_id}",
)
add_status(

2
ingestion/src/airflow_provider_openmetadata/lineage/config/loader.py
View File

@ -90,7 +90,7 @@ def get_lineage_config() -> AirflowLineageConfig:
if openmetadata_config_file:
with open(openmetadata_config_file, encoding="utf-8") as config_file:
config = json.load(config_file)
return AirflowLineageConfig.parse_obj(config)
return AirflowLineageConfig.model_validate(config)
# If nothing is configured, raise
raise ValueError("Missing lineage backend configuration")

8
ingestion/src/airflow_provider_openmetadata/lineage/runner.py
View File

@ -282,7 +282,7 @@ class AirflowLineageRunner:
for status in pipeline_status_list:
self.metadata.add_pipeline_status(
fqn=pipeline.fullyQualifiedName.__root__, status=status
fqn=pipeline.fullyQualifiedName.root, status=status
)
def add_lineage(self, pipeline: Pipeline, xlets: XLets) -> None:
@ -327,12 +327,12 @@ class AirflowLineageRunner:
else:
self.dag.log.warning(
f"Could not find [{to_xlet.entity.__name__}] [{to_xlet.fqn}] from "
f"[{pipeline.fullyQualifiedName.__root__}] outlets"
f"[{pipeline.fullyQualifiedName.root}] outlets"
)
else:
self.dag.log.warning(
f"Could not find [{from_xlet.entity.__name__}] [{from_xlet.fqn}] from "
f"[{pipeline.fullyQualifiedName.__root__}] inlets"
f"[{pipeline.fullyQualifiedName.root}] inlets"
)
def clean_lineage(self, pipeline: Pipeline, xlets: XLets):
@ -343,7 +343,7 @@ class AirflowLineageRunner:
"""
lineage_data = self.metadata.get_lineage_by_name(
entity=Pipeline,
fqn=pipeline.fullyQualifiedName.__root__,
fqn=pipeline.fullyQualifiedName.root,
up_depth=1,
down_depth=1,
)

4
ingestion/src/airflow_provider_openmetadata/lineage/status.py
View File

@ -91,7 +91,7 @@ def add_status(
task_status = []
# We will append based on the current registered status
if pipeline_status and pipeline_status.timestamp.__root__ == execution_date:
if pipeline_status and pipeline_status.timestamp.root == execution_date:
# If we are clearing a task, use the status of the new execution
task_status = [
task
@ -123,5 +123,5 @@ def add_status(
operator.log.info(f"Added status to DAG {updated_status}")
metadata.add_pipeline_status(
fqn=pipeline.fullyQualifiedName.__root__, status=updated_status
fqn=pipeline.fullyQualifiedName.root, status=updated_status
)

2
ingestion/src/metadata/automations/runner.py
View File

@ -37,7 +37,7 @@ def execute(encrypted_automation_workflow: AutomationWorkflow) -> Any:
)
automation_workflow = metadata.get_by_name(
entity=AutomationWorkflow, fqn=encrypted_automation_workflow.name.__root__
entity=AutomationWorkflow, fqn=encrypted_automation_workflow.name.root
)
return run_workflow(automation_workflow.request, automation_workflow, metadata)

6
ingestion/src/metadata/cli/lineage.py
View File

@ -31,8 +31,8 @@ logger = cli_logger()
class LineageWorkflow(BaseModel):
filePath: Optional[str]
query: Optional[str]
filePath: Optional[str] = None
query: Optional[str] = None
checkPatch: Optional[bool] = True
serviceName: str
workflowConfig: WorkflowConfig
@ -49,7 +49,7 @@ def run_lineage(config_path: Path) -> None:
config_dict = None
try:
config_dict = load_config_file(config_path)
workflow = LineageWorkflow.parse_obj(config_dict)
workflow = LineageWorkflow.model_validate(config_dict)
except Exception as exc:
logger.debug(traceback.format_exc())

8
ingestion/src/metadata/clients/aws_client.py
View File

@ -47,7 +47,7 @@ class AWSAssumeRoleException(Exception):
class AWSAssumeRoleCredentialWrapper(BaseModel):
accessKeyId: str
secretAccessKey: CustomSecretStr
sessionToken: Optional[str]
sessionToken: Optional[str] = None
class AWSClient:
@ -59,7 +59,7 @@ class AWSClient:
self.config = (
config
if isinstance(config, AWSCredentials)
else (AWSCredentials.parse_obj(config) if config else config)
else (AWSCredentials.model_validate(config) if config else config)
)
@staticmethod
@ -148,7 +148,7 @@ class AWSClient:
session = self.create_session()
if self.config.endPointURL is not None:
return session.client(
service_name=service_name, endpoint_url=self.config.endPointURL
service_name=service_name, endpoint_url=str(self.config.endPointURL)
)
return session.client(service_name=service_name)
@ -160,7 +160,7 @@ class AWSClient:
session = self.create_session()
if self.config.endPointURL is not None:
return session.resource(
service_name=service_name, endpoint_url=self.config.endPointURL
service_name=service_name, endpoint_url=str(self.config.endPointURL)
)
return session.resource(service_name=service_name)

2
ingestion/src/metadata/clients/azure_client.py
View File

@ -28,7 +28,7 @@ class AzureClient:
def __init__(self, credentials: "AzureCredentials"):
self.credentials = credentials
if not isinstance(credentials, AzureCredentials):
self.credentials = AzureCredentials.parse_obj(credentials)
self.credentials = AzureCredentials.model_validate(credentials)
def create_client(
self,

20
ingestion/src/metadata/clients/domo_client.py
View File

@ -64,10 +64,10 @@ class DomoDashboardDetails(DomoBaseModel):
Response from Domo API
"""
cardIds: Optional[List[int]]
collectionIds: Optional[List[int]]
description: Optional[str]
owners: Optional[List[DomoOwner]]
cardIds: Optional[List[int]] = None
collectionIds: Optional[List[int]] = None
description: Optional[str] = None
owners: Optional[List[DomoOwner]] = None
class DomoChartMetadataDetails(BaseModel):
@ -78,7 +78,7 @@ class DomoChartMetadataDetails(BaseModel):
class Config:
extra = Extra.allow
chartType: Optional[str]
chartType: Optional[str] = None
class DomoChartDetails(DomoBaseModel):
@ -87,7 +87,7 @@ class DomoChartDetails(DomoBaseModel):
"""
metadata: DomoChartMetadataDetails
description: Optional[str]
description: Optional[str] = None
class DomoClient:
@ -103,14 +103,10 @@ class DomoClient:
],
):
self.config = config
self.config.instanceDomain = (
self.config.instanceDomain[:-1]
if self.config.instanceDomain.endswith("/")
else self.config.instanceDomain
)
HEADERS.update({"X-DOMO-Developer-Token": self.config.accessToken})
client_config: ClientConfig = ClientConfig(
base_url=self.config.instanceDomain,
# AnyUrl string ends with / and the domo API does not respond properly if it has 2 // at the end
base_url=str(self.config.instanceDomain)[:-1],
api_version="api/",
auth_header="Authorization",
auth_token=lambda: ("no_token", 0),

2
ingestion/src/metadata/config/common.py
View File

@ -33,7 +33,7 @@ class DynamicTypedConfig(ConfigModel):
"""Class definition for Dynamic Typed Config"""
type: str
config: Optional[Any]
config: Optional[Any] = None
class WorkflowExecutionError(Exception):

6
ingestion/src/metadata/data_insight/processor/kpi/kpi_runner.py
View File

@ -65,8 +65,8 @@ class KpiRunner:
Kpi:
"""
start_date = entity.startDate.__root__
end_date = entity.endDate.__root__
start_date = entity.startDate.root
end_date = entity.endDate.root
if not start_date or not end_date:
logger.warning(
@ -128,7 +128,7 @@ class KpiRunner:
results = self.metadata.get_aggregated_data_insight_results(
start_ts=get_beginning_of_day_timestamp_mill(),
end_ts=get_end_of_day_timestamp_mill(),
data_insight_chart_nane=data_insight_chart_entity.name.__root__,
data_insight_chart_nane=data_insight_chart_entity.name.root,
data_report_index=data_insight_chart_entity.dataIndexType.value,
)
if results.data or tme.time() > timeout:

10
ingestion/src/metadata/data_insight/processor/kpi/run_result_registry.py
View File

@ -24,7 +24,7 @@ from metadata.generated.schema.dataInsight.type.percentageOfEntitiesWithDescript
from metadata.generated.schema.dataInsight.type.percentageOfEntitiesWithOwnerByType import (
PercentageOfEntitiesWithOwnerByType,
)
from metadata.generated.schema.type.basic import FullyQualifiedEntityName
from metadata.generated.schema.type.basic import FullyQualifiedEntityName, Timestamp
from metadata.utils.dispatch import enum_register
from metadata.utils.logger import profiler_interface_registry_logger
@ -81,13 +81,13 @@ def percentage_of_entities_with_description_kpi_result(
target_results.append(
KpiTarget(
name=target.name,
value=value,
value=str(value),
targetMet=value > ast.literal_eval(target.value),
)
)
return KpiResult(
timestamp=timestamp,
timestamp=Timestamp(timestamp),
targetResult=target_results,
kpiFqn=kpi_fqn,
)
@ -141,13 +141,13 @@ def percentage_of_entities_with_owner_kpi_result(
target_results.append(
KpiTarget(
name=target.name,
value=value,
value=str(value),
targetMet=value > ast.literal_eval(target.value),
)
)
return KpiResult(
timestamp=timestamp,
timestamp=Timestamp(timestamp),
targetResult=target_results,
kpiFqn=kpi_fqn,
)

2
ingestion/src/metadata/data_insight/processor/reports/cost_analysis_report_data_processor.py
View File

@ -237,7 +237,7 @@ class AggregatedCostAnalysisReportDataProcessor(DataProcessor):
days_before_timestamp = get_end_of_day_timestamp_mill(days=days)
if (
life_cycle.accessed
and life_cycle.accessed.timestamp.__root__ <= days_before_timestamp
and life_cycle.accessed.timestamp.root <= days_before_timestamp
):
data[UNUSED_DATA_ASSETS][COUNT][key] += 1
data[UNUSED_DATA_ASSETS][SIZE][key] += size or 0

3
ingestion/src/metadata/data_insight/processor/reports/data_processor.py
View File

@ -20,6 +20,7 @@ from datetime import datetime, timezone
from typing import Callable, Iterable, Optional
from metadata.generated.schema.analytics.reportData import ReportData
from metadata.generated.schema.type.basic import Timestamp
from metadata.ingestion.api.status import Status
from metadata.ingestion.ometa.ometa_api import OpenMetadata
@ -42,7 +43,7 @@ class DataProcessor(abc.ABC):
def __init__(self, metadata: OpenMetadata):
self.metadata = metadata
self.timestamp = datetime.now(timezone.utc).timestamp() * 1000
self.timestamp = Timestamp(int(datetime.now(timezone.utc).timestamp() * 1000))
self.processor_status = Status()
self._refined_data = {}
self.post_hook: Optional[Callable] = None

12
ingestion/src/metadata/data_insight/processor/reports/entity_report_data_processor.py
View File

@ -95,7 +95,7 @@ class EntityReportDataProcessor(DataProcessor):
return None
if isinstance(owner, EntityReferenceList):
return owner.__root__[0].name
return owner.root[0].name
if owner.type == "team":
return owner.name
@ -113,7 +113,7 @@ class EntityReportDataProcessor(DataProcessor):
teams = entity_reference.teams
if teams:
return teams.__root__[0].name # We'll return the first team listed
return teams.root[0].name # We'll return the first team listed
return None
@ -136,7 +136,7 @@ class EntityReportDataProcessor(DataProcessor):
return True
if entity.description and not entity.description.__root__ == "":
if entity.description and not entity.description.root == "":
return True
return False
@ -163,7 +163,7 @@ class EntityReportDataProcessor(DataProcessor):
yield ReportData(
timestamp=self.timestamp,
reportDataType=ReportDataType.entityReportData.value,
data=EntityReportData.parse_obj(data),
data=EntityReportData.model_validate(data),
) # type: ignore
def refine(self, entity: Type[T]) -> None:
@ -195,7 +195,7 @@ class EntityReportDataProcessor(DataProcessor):
except Exception:
self.processor_status.failed(
StackTraceError(
name=entity.name.__root__,
name=entity.name.root,
error="Error retrieving team",
stackTrace=traceback.format_exc(),
)
@ -256,7 +256,7 @@ class EntityReportDataProcessor(DataProcessor):
str(team)
][str(entity_tier)].update(data_blob_for_entity_counter)
self.processor_status.scanned(entity.name.__root__)
self.processor_status.scanned(entity.name.root)
def get_status(self):
return self.processor_status

36
ingestion/src/metadata/data_insight/processor/reports/web_analytic_report_data_processor.py
View File

@ -101,7 +101,7 @@ class WebAnalyticEntityViewReportDataProcessor(DataProcessor):
while True:
event = yield refined_data
split_url = [url for url in event.eventData.url.__root__.split("/") if url] # type: ignore
split_url = [url for url in event.eventData.url.root.split("/") if url] # type: ignore
if not split_url or split_url[0] not in ENTITIES:
continue
@ -120,7 +120,7 @@ class WebAnalyticEntityViewReportDataProcessor(DataProcessor):
# the URL we'll try again from the new event.
try:
entity_href = re.search(
re_pattern, event.eventData.fullUrl.__root__
re_pattern, event.eventData.fullUrl.root
).group(1)
refined_data[entity_obj.fqn]["entityHref"] = entity_href
except IndexError:
@ -145,7 +145,7 @@ class WebAnalyticEntityViewReportDataProcessor(DataProcessor):
try:
tags = (
[tag.tagFQN.__root__ for tag in entity.tags]
[tag.tagFQN.root for tag in entity.tags]
if entity.tags
else None
)
@ -159,7 +159,7 @@ class WebAnalyticEntityViewReportDataProcessor(DataProcessor):
try:
owner = entity.owner.name if entity.owner else None
owner_id = str(entity.owner.id.__root__) if entity.owner else None
owner_id = str(entity.owner.id.root) if entity.owner else None
except AttributeError as exc:
owner = None
owner_id = None
@ -173,7 +173,7 @@ class WebAnalyticEntityViewReportDataProcessor(DataProcessor):
try:
entity_href = re.search(
re_pattern, event.eventData.fullUrl.__root__
re_pattern, event.eventData.fullUrl.root
).group(1)
except IndexError:
entity_href = None
@ -181,7 +181,7 @@ class WebAnalyticEntityViewReportDataProcessor(DataProcessor):
if (
owner_id is not None
and event.eventData is not None
and owner_id == str(event.eventData.userId.__root__)
and owner_id == str(event.eventData.userId.root)
): # type: ignore
# we won't count views if the owner is the one visiting
# the entity
@ -208,7 +208,7 @@ class WebAnalyticEntityViewReportDataProcessor(DataProcessor):
yield ReportData(
timestamp=self.timestamp,
reportDataType=ReportDataType.webAnalyticEntityViewReportData.value,
data=WebAnalyticEntityViewReportData.parse_obj(
data=WebAnalyticEntityViewReportData.model_validate(
self._refined_data[data]
),
) # type: ignore
@ -273,7 +273,7 @@ class WebAnalyticUserActivityReportDataProcessor(DataProcessor):
return {
"totalSessions": total_sessions,
"totalSessionDuration": total_session_duration_seconds,
"totalSessionDuration": int(total_session_duration_seconds),
}
def _get_user_details(self, user_id: str) -> dict:
@ -298,12 +298,12 @@ class WebAnalyticUserActivityReportDataProcessor(DataProcessor):
teams = user_entity.teams
return {
"user_name": user_entity.name.__root__,
"team": teams.__root__[0].name if teams else None,
"user_name": user_entity.name.root,
"team": teams.root[0].name if teams else None,
}
def _refine_user_event(self) -> Generator[dict, WebAnalyticEventData, None]:
"""Corountine to process user event from web analytic event
"""Coroutine to process user event from web analytic event
Yields:
Generator[dict, WebAnalyticEventData, None]: _description_
@ -313,9 +313,9 @@ class WebAnalyticUserActivityReportDataProcessor(DataProcessor):
while True:
event = yield self._refined_data
user_id = str(event.eventData.userId.__root__) # type: ignore
session_id = str(event.eventData.sessionId.__root__) # type: ignore
timestamp = event.timestamp.__root__ # type: ignore
user_id = str(event.eventData.userId.root) # type: ignore
session_id = str(event.eventData.sessionId.root) # type: ignore
timestamp = event.timestamp.root # type: ignore
if not user_details.get(user_id):
user_details_data = self._get_user_details(user_id)
@ -351,8 +351,7 @@ class WebAnalyticUserActivityReportDataProcessor(DataProcessor):
def fetch_data(self) -> Iterable[WebAnalyticEventData]:
if CACHED_EVENTS:
for event in CACHED_EVENTS:
yield event
yield from CACHED_EVENTS
else:
CACHED_EVENTS.extend(
self.metadata.list_entities(
@ -364,8 +363,7 @@ class WebAnalyticUserActivityReportDataProcessor(DataProcessor):
},
).entities
)
for event in CACHED_EVENTS:
yield event
yield from CACHED_EVENTS
def yield_refined_data(self) -> Iterable[ReportData]:
"""Yield refined data"""
@ -373,7 +371,7 @@ class WebAnalyticUserActivityReportDataProcessor(DataProcessor):
yield ReportData(
timestamp=self.timestamp,
reportDataType=ReportDataType.webAnalyticUserActivityReportData.value,
data=WebAnalyticUserActivityReportData.parse_obj(
data=WebAnalyticUserActivityReportData.model_validate(
self._refined_data[user_id]
),
) # type: ignore

8
ingestion/src/metadata/data_insight/producer/cost_analysis_producer.py
View File

@ -34,8 +34,8 @@ class CostAnalysisReportData(BaseModel):
"""
entity: Entity
life_cycle: Optional[LifeCycle]
size: Optional[float]
life_cycle: Optional[LifeCycle] = None
size: Optional[float] = None
class CostAnalysisProducer(ProducerInterface):
@ -46,9 +46,9 @@ class CostAnalysisProducer(ProducerInterface):
) -> bool:
return (
hasattr(database_service.connection.config, "supportsUsageExtraction")
and database_service.connection.config.supportsUsageExtraction.__root__
and database_service.connection.config.supportsUsageExtraction.root
and hasattr(database_service.connection.config, "supportsProfiler")
and database_service.connection.config.supportsProfiler.__root__
and database_service.connection.config.supportsProfiler.root
)
def _check_life_cycle_and_size_data(

3
ingestion/src/metadata/data_insight/producer/web_analytics_producer.py
View File

@ -80,5 +80,4 @@ class WebAnalyticsProducer(ProducerInterface):
"""fetch data for web analytics event"""
events = self._get_events(None, limit, fields)
for entity in events.entities:
yield entity
yield from events.entities

2
ingestion/src/metadata/data_quality/api/models.py
View File

@ -35,7 +35,7 @@ class TestCaseDefinition(ConfigModel):
description: Optional[str] = None
testDefinitionName: str
columnName: Optional[str] = None
parameterValues: Optional[List[TestCaseParameterValue]]
parameterValues: Optional[List[TestCaseParameterValue]] = None
computePassedFailedRowCount: Optional[bool] = False

30
ingestion/src/metadata/data_quality/processor/test_case_runner.py
View File

@ -60,8 +60,8 @@ class TestCaseRunner(Processor):
self.metadata = metadata
self.processor_config: TestSuiteProcessorConfig = (
TestSuiteProcessorConfig.parse_obj(
self.config.processor.dict().get("config")
TestSuiteProcessorConfig.model_validate(
self.config.processor.model_dump().get("config")
)
)
@ -82,16 +82,16 @@ class TestCaseRunner(Processor):
test_suite_fqn=fqn.build(
None,
TestSuite,
table_fqn=record.table.fullyQualifiedName.__root__,
table_fqn=record.table.fullyQualifiedName.root,
),
table_fqn=record.table.fullyQualifiedName.__root__,
table_fqn=record.table.fullyQualifiedName.root,
)
if not test_cases:
return Either(
left=StackTraceError(
name="No test Cases",
error=f"No tests cases found for table {record.table.fullyQualifiedName.__root__}",
error=f"No tests cases found for table {record.table.fullyQualifiedName.root}",
)
)
@ -162,9 +162,7 @@ class TestCaseRunner(Processor):
return test_cases
test_cases = deepcopy(test_cases) or []
test_case_names = (
{test_case.name.__root__ for test_case in test_cases}
if test_cases
else set()
{test_case.name.root for test_case in test_cases} if test_cases else set()
)
# we'll check the test cases defined in the CLI config file and not present in the platform
@ -196,10 +194,10 @@ class TestCaseRunner(Processor):
description=test_case_to_create.description,
displayName=test_case_to_create.displayName,
testDefinition=FullyQualifiedEntityName(
__root__=test_case_to_create.testDefinitionName
test_case_to_create.testDefinitionName
),
entityLink=EntityLink(
__root__=entity_link.get_entity_link(
entity_link.get_entity_link(
Table,
fqn=table_fqn,
column_name=test_case_to_create.columnName,
@ -245,11 +243,11 @@ class TestCaseRunner(Processor):
test_case_to_update.name for test_case_to_update in test_cases_to_update
}
for indx, test_case in enumerate(deepcopy(test_cases)):
if test_case.name.__root__ in test_cases_to_update_names:
if test_case.name.root in test_cases_to_update_names:
test_case_definition = next(
test_case_to_update
for test_case_to_update in test_cases_to_update
if test_case_to_update.name == test_case.name.__root__
if test_case_to_update.name == test_case.name.root
)
updated_test_case = self.metadata.patch_test_case_definition(
test_case=test_case,
@ -281,7 +279,7 @@ class TestCaseRunner(Processor):
)
if TestPlatform.OpenMetadata not in test_definition.testPlatforms:
logger.debug(
f"Test case {test_case.name.__root__} is not an OpenMetadata test case."
f"Test case {test_case.name.root} is not an OpenMetadata test case."
)
continue
om_test_cases.append(test_case)
@ -294,15 +292,15 @@ class TestCaseRunner(Processor):
"""Execute the test case and return the result, if any"""
try:
test_result = test_suite_runner.run_and_handle(test_case)
self.status.scanned(test_case.fullyQualifiedName.__root__)
self.status.scanned(test_case.fullyQualifiedName.root)
return test_result
except Exception as exc:
error = f"Could not run test case {test_case.name.__root__}: {exc}"
error = f"Could not run test case {test_case.name.root}: {exc}"
logger.debug(traceback.format_exc())
logger.error(error)
self.status.failed(
StackTraceError(
name=test_case.name.__root__,
name=test_case.name.root,
error=error,
stackTrace=traceback.format_exc(),
)

2
ingestion/src/metadata/data_quality/runner/base_test_suite_source.py
View File

@ -71,7 +71,7 @@ class BaseTestSuiteRunner:
DatabaseService.__config__
"""
config_copy = deepcopy(
config.source.serviceConnection.__root__.config # type: ignore
config.source.serviceConnection.root.config # type: ignore
)
if hasattr(
config_copy, # type: ignore

4
ingestion/src/metadata/data_quality/runner/core.py
View File

@ -31,8 +31,8 @@ class DataTestsRunner:
def run_and_handle(self, test_case: TestCase):
"""run and handle test case validation"""
logger.info(
f"Executing test case {test_case.name.__root__} "
f"for entity {self.test_runner_interface.table_entity.fullyQualifiedName.__root__}"
f"Executing test case {test_case.name.root} "
f"for entity {self.test_runner_interface.table_entity.fullyQualifiedName.root}"
)
test_result = self.test_runner_interface.run_test_case(
test_case,

20
ingestion/src/metadata/data_quality/source/test_suite.py
View File

@ -72,7 +72,7 @@ class TestSuiteSource(Source):
"""
table: Table = self.metadata.get_by_name(
entity=Table,
fqn=self.source_config.entityFullyQualifiedName.__root__,
fqn=self.source_config.entityFullyQualifiedName.root,
fields=["tableProfilerConfig", "testSuite"],
)
@ -86,7 +86,7 @@ class TestSuiteSource(Source):
test_cases = self.metadata.list_all_entities(
entity=TestCase,
fields=["testSuite", "entityLink", "testDefinition"],
params={"testSuiteId": test_suite.id.__root__},
params={"testSuiteId": test_suite.id.root},
)
test_cases = cast(List[TestCase], test_cases) # satisfy type checker
@ -110,7 +110,7 @@ class TestSuiteSource(Source):
yield Either(
left=StackTraceError(
name="Missing Table",
error=f"Could not retrieve table entity for {self.source_config.entityFullyQualifiedName.__root__}."
error=f"Could not retrieve table entity for {self.source_config.entityFullyQualifiedName.root}."
" Make sure the table exists in OpenMetadata and/or the JWT Token provided is valid.",
)
)
@ -125,31 +125,31 @@ class TestSuiteSource(Source):
name=fqn.build(
None,
TestSuite,
table_fqn=self.source_config.entityFullyQualifiedName.__root__,
table_fqn=self.source_config.entityFullyQualifiedName.root,
),
displayName=f"{self.source_config.entityFullyQualifiedName.__root__} Test Suite",
displayName=f"{self.source_config.entityFullyQualifiedName.root} Test Suite",
description="Test Suite created from YAML processor config file",
owner=None,
executableEntityReference=self.source_config.entityFullyQualifiedName.__root__,
executableEntityReference=self.source_config.entityFullyQualifiedName.root,
)
yield Either(
right=TableAndTests(
executable_test_suite=executable_test_suite,
service_type=self.config.source.serviceConnection.__root__.config.type.value,
service_type=self.config.source.serviceConnection.root.config.type.value,
)
)
test_suite: Optional[TestSuite] = None
if table.testSuite:
test_suite = self.metadata.get_by_id(
entity=TestSuite, entity_id=table.testSuite.id.__root__
entity=TestSuite, entity_id=table.testSuite.id.root
)
if test_suite and not test_suite.executable:
yield Either(
left=StackTraceError(
name="Non-executable Test Suite",
error=f"The table {self.source_config.entityFullyQualifiedName.__root__} "
error=f"The table {self.source_config.entityFullyQualifiedName.root} "
"has a test suite that is not executable.",
)
)
@ -161,7 +161,7 @@ class TestSuiteSource(Source):
right=TableAndTests(
table=table,
test_cases=test_suite_cases,
service_type=self.config.source.serviceConnection.__root__.config.type.value,
service_type=self.config.source.serviceConnection.root.config.type.value,
)
)

2
ingestion/src/metadata/data_quality/validations/column/base/columnValuesToBeNotInSet.py
View File

@ -59,7 +59,7 @@ class BaseColumnValuesToBeNotInSetValidator(BaseTestValidator):
except (ValueError, RuntimeError) as exc:
msg = (
f"Error computing {self.test_case.name} for "
f"{get_table_fqn(self.test_case.entityLink.__root__)}: {exc}"
f"{get_table_fqn(self.test_case.entityLink.root)}: {exc}"
)
logger.debug(traceback.format_exc())
logger.warning(msg)

2
ingestion/src/metadata/data_quality/validations/column/pandas/columnValueLengthsToBeBetween.py
View File

@ -38,7 +38,7 @@ class ColumnValueLengthsToBeBetweenValidator(
SQALikeColumn:
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
self.runner,
)

2
ingestion/src/metadata/data_quality/validations/column/pandas/columnValueMaxToBeBetween.py
View File

@ -37,7 +37,7 @@ class ColumnValueMaxToBeBetweenValidator(
SQALikeColumn: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
self.runner,
)

2
ingestion/src/metadata/data_quality/validations/column/pandas/columnValueMeanToBeBetween.py
View File

@ -38,7 +38,7 @@ class ColumnValueMeanToBeBetweenValidator(
SQALikeColumn: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
self.runner,
)

2
ingestion/src/metadata/data_quality/validations/column/pandas/columnValueMedianToBeBetween.py
View File

@ -37,7 +37,7 @@ class ColumnValueMedianToBeBetweenValidator(
SQALikeColumn: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
self.runner,
)

2
ingestion/src/metadata/data_quality/validations/column/pandas/columnValueMinToBeBetween.py
View File

@ -37,7 +37,7 @@ class ColumnValueMinToBeBetweenValidator(
SQALikeColumn: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
self.runner,
)

2
ingestion/src/metadata/data_quality/validations/column/pandas/columnValueStdDevToBeBetween.py
View File

@ -37,7 +37,7 @@ class ColumnValueStdDevToBeBetweenValidator(
SQALikeColumn: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
self.runner,
)

2
ingestion/src/metadata/data_quality/validations/column/pandas/columnValuesMissingCount.py
View File

@ -37,7 +37,7 @@ class ColumnValuesMissingCountValidator(
SQALikeColumn: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
self.runner,
)

2
ingestion/src/metadata/data_quality/validations/column/pandas/columnValuesSumToBeBetween.py
View File

@ -37,7 +37,7 @@ class ColumnValuesSumToBeBetweenValidator(
SQALikeColumn: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
self.runner,
)

2
ingestion/src/metadata/data_quality/validations/column/pandas/columnValuesToBeBetween.py
View File

@ -37,7 +37,7 @@ class ColumnValuesToBeBetweenValidator(
SQALikeColumn: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
self.runner,
)

2
ingestion/src/metadata/data_quality/validations/column/pandas/columnValuesToBeInSet.py
View File

@ -40,7 +40,7 @@ class ColumnValuesToBeInSetValidator(
SQALikeColumn: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
self.runner,
)

2
ingestion/src/metadata/data_quality/validations/column/pandas/columnValuesToBeNotInSet.py
View File

@ -37,7 +37,7 @@ class ColumnValuesToBeNotInSetValidator(
SQALikeColumn: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
self.runner,
)

2
ingestion/src/metadata/data_quality/validations/column/pandas/columnValuesToBeNotNull.py
View File

@ -37,7 +37,7 @@ class ColumnValuesToBeNotNullValidator(
SQALikeColumn: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
self.runner,
)

2
ingestion/src/metadata/data_quality/validations/column/pandas/columnValuesToBeUnique.py
View File

@ -37,7 +37,7 @@ class ColumnValuesToBeUniqueValidator(
SQALikeColumn: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
self.runner,
)

2
ingestion/src/metadata/data_quality/validations/column/pandas/columnValuesToMatchRegex.py
View File

@ -37,7 +37,7 @@ class ColumnValuesToMatchRegexValidator(
SQALikeColumn: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
self.runner,
)

2
ingestion/src/metadata/data_quality/validations/column/pandas/columnValuesToNotMatchRegex.py
View File

@ -37,7 +37,7 @@ class ColumnValuesToNotMatchRegexValidator(
SQALikeColumn: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
self.runner,
)

2
ingestion/src/metadata/data_quality/validations/column/sqlalchemy/columnValueLengthsToBeBetween.py
View File

@ -40,7 +40,7 @@ class ColumnValueLengthsToBeBetweenValidator(
Column: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
inspect(self.runner.table).c,
)

2
ingestion/src/metadata/data_quality/validations/column/sqlalchemy/columnValueMaxToBeBetween.py
View File

@ -37,7 +37,7 @@ class ColumnValueMaxToBeBetweenValidator(
Column: _description_
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
inspect(self.runner.table).c,
)

2
ingestion/src/metadata/data_quality/validations/column/sqlalchemy/columnValueMeanToBeBetween.py
View File

@ -38,7 +38,7 @@ class ColumnValueMeanToBeBetweenValidator(
Column: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
inspect(self.runner.table).c,
)

2
ingestion/src/metadata/data_quality/validations/column/sqlalchemy/columnValueMedianToBeBetween.py
View File

@ -38,7 +38,7 @@ class ColumnValueMedianToBeBetweenValidator(
Column: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
inspect(self.runner.table).c,
)

2
ingestion/src/metadata/data_quality/validations/column/sqlalchemy/columnValueMinToBeBetween.py
View File

@ -38,7 +38,7 @@ class ColumnValueMinToBeBetweenValidator(
Column: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
inspect(self.runner.table).c,
)

2
ingestion/src/metadata/data_quality/validations/column/sqlalchemy/columnValueStdDevToBeBetween.py
View File

@ -38,7 +38,7 @@ class ColumnValueStdDevToBeBetweenValidator(
Column: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
inspect(self.runner.table).c,
)

2
ingestion/src/metadata/data_quality/validations/column/sqlalchemy/columnValuesMissingCount.py
View File

@ -41,7 +41,7 @@ class ColumnValuesMissingCountValidator(
SQALikeColumn: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
inspect(self.runner.table).c,
)

2
ingestion/src/metadata/data_quality/validations/column/sqlalchemy/columnValuesSumToBeBetween.py
View File

@ -38,7 +38,7 @@ class ColumnValuesSumToBeBetweenValidator(
Column: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
inspect(self.runner.table).c,
)

2
ingestion/src/metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeBetween.py
View File

@ -38,7 +38,7 @@ class ColumnValuesToBeBetweenValidator(
Column: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
inspect(self.runner.table).c,
)

2
ingestion/src/metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeInSet.py
View File

@ -38,7 +38,7 @@ class ColumnValuesToBeInSetValidator(
Column: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
inspect(self.runner.table).c,
)

2
ingestion/src/metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeNotInSet.py
View File

@ -38,7 +38,7 @@ class ColumnValuesToBeNotInSetValidator(
Column: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
inspect(self.runner.table).c,
)

2
ingestion/src/metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeNotNull.py
View File

@ -41,7 +41,7 @@ class ColumnValuesToBeNotNullValidator(
Column: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
inspect(self.runner.table).c,
)

2
ingestion/src/metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeUnique.py
View File

@ -40,7 +40,7 @@ class ColumnValuesToBeUniqueValidator(
Column: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
inspect(self.runner.table).c,
)

2
ingestion/src/metadata/data_quality/validations/column/sqlalchemy/columnValuesToMatchRegex.py
View File

@ -42,7 +42,7 @@ class ColumnValuesToMatchRegexValidator(
Column: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
inspect(self.runner.table).c,
)

2
ingestion/src/metadata/data_quality/validations/column/sqlalchemy/columnValuesToNotMatchRegex.py
View File

@ -42,7 +42,7 @@ class ColumnValuesToNotMatchRegexValidator(
SQALikeColumn: column
"""
return self.get_column_name(
self.test_case.entityLink.__root__,
self.test_case.entityLink.root,
inspect(self.runner.table).c,
)

2
ingestion/src/metadata/examples/workflows/bigquery_profiler.yaml
View File

@ -11,7 +11,7 @@ source:
privateKeyId: privateKeyID
privateKey: "-----BEGIN PRIVATE KEY-----\nmySuperSecurePrivateKey==\n-----END PRIVATE KEY-----\n"
clientEmail: client@email.secure
clientId: 1234567890
clientId: "1234567890"
authUri: https://accounts.google.com/o/oauth2/auth
tokenUri: https://oauth2.googleapis.com/token
authProviderX509CertUrl: https://www.googleapis.com/oauth2/v1/certs

26
ingestion/src/metadata/great_expectations/action.py
View File

@ -32,6 +32,8 @@ from great_expectations.core.expectation_validation_result import (
from great_expectations.data_asset.data_asset import DataAsset
from great_expectations.data_context.data_context import DataContext
from metadata.generated.schema.type.basic import Timestamp
try:
from great_expectations.data_context.types.resource_identifiers import (
GeCloudIdentifier, # type: ignore
@ -219,7 +221,7 @@ class OpenMetadataValidationAction(ValidationAction):
entity=Table, fields=["testSuite"]
).entities
if f"{database}.{schema_name}.{table_name}"
in entity.fullyQualifiedName.__root__
in entity.fullyQualifiedName.root
]
if len(table_entity) > 1:
@ -248,14 +250,14 @@ class OpenMetadataValidationAction(ValidationAction):
if table_entity.testSuite:
test_suite = self.ometa_conn.get_by_name(
TestSuite, table_entity.testSuite.fullyQualifiedName.__root__
TestSuite, table_entity.testSuite.fullyQualifiedName.root
)
test_suite = cast(TestSuite, test_suite)
return test_suite
create_test_suite = CreateTestSuiteRequest(
name=f"{table_entity.fullyQualifiedName.__root__}.TestSuite",
executableEntityReference=table_entity.fullyQualifiedName.__root__,
name=f"{table_entity.fullyQualifiedName.root}.TestSuite",
executableEntityReference=table_entity.fullyQualifiedName.root,
) # type: ignore
test_suite = self.ometa_conn.create_or_update_executable_test_suite(
create_test_suite
@ -403,7 +405,7 @@ class OpenMetadataValidationAction(ValidationAction):
)
test_case_fqn = self._build_test_case_fqn(
table_entity.fullyQualifiedName.__root__,
table_entity.fullyQualifiedName.root,
result,
)
@ -411,27 +413,29 @@ class OpenMetadataValidationAction(ValidationAction):
test_case_fqn,
entity_link=get_entity_link(
Table,
fqn=table_entity.fullyQualifiedName.__root__,
fqn=table_entity.fullyQualifiedName.root,
column_name=fqn.split_test_case_fqn(test_case_fqn).column,
),
test_suite_fqn=test_suite.fullyQualifiedName.__root__,
test_definition_fqn=test_definition.fullyQualifiedName.__root__,
test_suite_fqn=test_suite.fullyQualifiedName.root,
test_definition_fqn=test_definition.fullyQualifiedName.root,
test_case_parameter_values=self._get_test_case_params_value(result),
)
self.ometa_conn.add_test_case_results(
test_results=TestCaseResult(
timestamp=int(datetime.now(timezone.utc).timestamp() * 1000),
timestamp=Timestamp(
int(datetime.now(tz=timezone.utc).timestamp() * 1000)
),
testCaseStatus=TestCaseStatus.Success
if result["success"]
else TestCaseStatus.Failed,
testResultValue=self._get_test_result_value(result),
), # type: ignore
test_case_fqn=test_case.fullyQualifiedName.__root__,
test_case_fqn=test_case.fullyQualifiedName.root,
)
logger.debug(
f"Test case result for {test_case.fullyQualifiedName.__root__} successfully ingested"
f"Test case result for {test_case.fullyQualifiedName.root} successfully ingested"
)
except Exception as exc:

2
ingestion/src/metadata/great_expectations/utils/ometa_config_handler.py
View File

@ -86,4 +86,4 @@ def render_template(environment: Environment, template_file: str = "config.yml")
def create_ometa_connection_obj(config: str) -> OpenMetadataConnection:
"""Create OpenMetadata connection"""
return OpenMetadataConnection.parse_obj(yaml.safe_load(config))
return OpenMetadataConnection.model_validate(yaml.safe_load(config))

2
ingestion/src/metadata/ingestion/api/common.py
View File

@ -33,7 +33,7 @@ class ConfigModel(BaseModel):
class DynamicTypedConfig(ConfigModel):
type: str
config: Optional[Any]
config: Optional[Any] = None
class WorkflowExecutionError(Exception):

2
ingestion/src/metadata/ingestion/api/delete.py
View File

@ -43,7 +43,7 @@ def delete_entity_from_source(
try:
entity_state = metadata.list_all_entities(entity=entity_type, params=params)
for entity in entity_state:
if str(entity.fullyQualifiedName.__root__) not in entity_source_state:
if str(entity.fullyQualifiedName.root) not in entity_source_state:
yield Either(
right=DeleteEntity(
entity=entity,

14
ingestion/src/metadata/ingestion/api/models.py
View File

@ -13,7 +13,7 @@ Generic models
"""
from typing import Generic, Optional, TypeVar
from pydantic import BaseModel
from pydantic import BaseModel, Field
from metadata.generated.schema.entity.services.ingestionPipelines.status import (
StackTraceError,
@ -25,11 +25,9 @@ T = TypeVar("T")
class Either(BaseModel, Generic[T]):
"""
Any execution should return us Either an Entity of an error for us to handle
- left: Optional error we encounter
- right: Correct instance of an Entity
"""
"""Any execution should return us Either an Entity of an error for us to handle"""
left: Optional[StackTraceError]
right: Optional[T]
left: Optional[StackTraceError] = Field(
None, description="Error encountered during execution"
)
right: Optional[T] = Field(None, description="Correct instance of an Entity")

26
ingestion/src/metadata/ingestion/api/parser.py
View File

@ -261,7 +261,7 @@ def _parse_validation_err(validation_error: ValidationError) -> str:
if len(err.get("loc")) == 1
else f"Extra parameter in {err.get('loc')}"
for err in validation_error.errors()
if err.get("type") == "value_error.extra"
if err.get("type") == "extra_forbidden"
]
extra_fields = [
@ -269,7 +269,7 @@ def _parse_validation_err(validation_error: ValidationError) -> str:
if len(err.get("loc")) == 1
else f"Missing parameter in {err.get('loc')}"
for err in validation_error.errors()
if err.get("type") == "value_error.missing"
if err.get("type") == "missing"
]
invalid_fields = [
@ -277,7 +277,7 @@ def _parse_validation_err(validation_error: ValidationError) -> str:
if len(err.get("loc")) == 1
else f"Invalid parameter value for {err.get('loc')}"
for err in validation_error.errors()
if err.get("type") not in ("value_error.missing", "value_error.extra")
if err.get("type") not in ("missing", "extra")
]
return "\t - " + "\n\t - ".join(missing_fields + extra_fields + invalid_fields)
@ -291,7 +291,7 @@ def _unsafe_parse_config(config: dict, cls: Type[T], message: str) -> None:
logger.debug(f"Parsing message: [{message}]")
# Parse the service connection dictionary with the scoped class
try:
cls.parse_obj(config)
cls.model_validate(config)
except ValidationError as err:
logger.debug(
f"The supported properties for {cls.__name__} are {list(cls.__fields__.keys())}"
@ -309,10 +309,10 @@ def _unsafe_parse_dbt_config(config: dict, cls: Type[T], message: str) -> None:
# Parse the oneOf config types of dbt to check
dbt_config_type = config["dbtConfigSource"]["dbtConfigType"]
dbt_config_class = DBT_CONFIG_TYPE_MAP.get(dbt_config_type)
dbt_config_class.parse_obj(config["dbtConfigSource"])
dbt_config_class.model_validate(config["dbtConfigSource"])
# Parse the entire dbtPipeline object
cls.parse_obj(config)
cls.model_validate(config)
except ValidationError as err:
logger.debug(
f"The supported properties for {cls.__name__} are {list(cls.__fields__.keys())}"
@ -437,21 +437,17 @@ def parse_workflow_config_gracefully(
"""
try:
workflow_config = OpenMetadataWorkflowConfig.parse_obj(config_dict)
workflow_config = OpenMetadataWorkflowConfig.model_validate(config_dict)
return workflow_config
except ValidationError as original_error:
try:
parse_workflow_source(config_dict)
WorkflowConfig.parse_obj(config_dict["workflowConfig"])
WorkflowConfig.model_validate(config_dict["workflowConfig"])
except (ValidationError, InvalidWorkflowException) as scoped_error:
if isinstance(scoped_error, ValidationError):
# Let's catch validations of internal Workflow models, not the Workflow itself
object_error = (
scoped_error.model.__name__
if scoped_error.model is not None
else "workflow"
)
object_error = scoped_error.title or "workflow"
raise ParsingConfigurationError(
f"We encountered an error parsing the configuration of your {object_error}.\n"
"You might need to review your config based on the original cause of this failure:\n"
@ -483,7 +479,7 @@ def parse_ingestion_pipeline_config_gracefully(
"""
try:
ingestion_pipeline = IngestionPipeline.parse_obj(config_dict)
ingestion_pipeline = IngestionPipeline.model_validate(config_dict)
return ingestion_pipeline
except ValidationError:
@ -518,7 +514,7 @@ def parse_automation_workflow_gracefully(
"""
try:
automation_workflow = AutomationWorkflow.parse_obj(config_dict)
automation_workflow = AutomationWorkflow.model_validate(config_dict)
return automation_workflow
except ValidationError:

19
ingestion/src/metadata/ingestion/api/status.py
View File

@ -16,6 +16,7 @@ import time
from typing import Any, Dict, List
from pydantic import BaseModel, Field
from typing_extensions import Annotated
from metadata.generated.schema.entity.services.ingestionPipelines.status import (
StackTraceError,
@ -31,17 +32,15 @@ class Status(BaseModel):
Class to handle status
"""
source_start_time: Any
source_start_time: float = Field(
default_factory=lambda: time.time() # pylint: disable=unnecessary-lambda
)
records: List[Any] = Field(default_factory=list)
updated_records: List[Any] = Field(default_factory=list)
warnings: List[Any] = Field(default_factory=list)
filtered: List[Dict[str, str]] = Field(default_factory=list)
failures: List[StackTraceError] = Field(default_factory=list)
def __init__(self, **data):
super().__init__(**data)
self.source_start_time = time.time()
records: Annotated[List[Any], Field(default_factory=list)]
updated_records: Annotated[List[Any], Field(default_factory=list)]
warnings: Annotated[List[Any], Field(default_factory=list)]
filtered: Annotated[List[Dict[str, str]], Field(default_factory=list)]
failures: Annotated[List[StackTraceError], Field(default_factory=list)]
def scanned(self, record: Any) -> None:
"""

34
ingestion/src/metadata/ingestion/bulksink/metadata_usage.py
View File

@ -21,7 +21,7 @@ import json
import os
import shutil
import traceback
from datetime import datetime
from datetime import datetime, timezone
from pathlib import Path
from typing import List, Optional
@ -99,7 +99,7 @@ class MetadataUsageBulkSink(BulkSink):
metadata: OpenMetadata,
pipeline_name: Optional[str] = None,
):
config = MetadataUsageSinkConfig.parse_obj(config_dict)
config = MetadataUsageSinkConfig.model_validate(config_dict)
return cls(config, metadata)
def __populate_table_usage_map(
@ -109,8 +109,8 @@ class MetadataUsageBulkSink(BulkSink):
Method Either initialise the map data or
update existing data with information from new queries on the same table
"""
if not self.table_usage_map.get(table_entity.id.__root__):
self.table_usage_map[table_entity.id.__root__] = {
if not self.table_usage_map.get(table_entity.id.root):
self.table_usage_map[table_entity.id.root] = {
"table_entity": table_entity,
"usage_count": table_usage.count,
"usage_date": table_usage.date,
@ -118,7 +118,7 @@ class MetadataUsageBulkSink(BulkSink):
"database_schema": table_usage.databaseSchema,
}
else:
self.table_usage_map[table_entity.id.__root__][
self.table_usage_map[table_entity.id.root][
"usage_count"
] += table_usage.count
@ -139,10 +139,10 @@ class MetadataUsageBulkSink(BulkSink):
value_dict["table_entity"], table_usage_request
)
logger.info(
f"Successfully table usage published for {value_dict['table_entity'].fullyQualifiedName.__root__}"
f"Successfully table usage published for {value_dict['table_entity'].fullyQualifiedName.root}"
)
self.status.scanned(
f"Table: {value_dict['table_entity'].fullyQualifiedName.__root__}"
f"Table: {value_dict['table_entity'].fullyQualifiedName.root}"
)
except ValidationError as err:
logger.debug(traceback.format_exc())
@ -150,13 +150,13 @@ class MetadataUsageBulkSink(BulkSink):
f"Cannot construct UsageRequest from {value_dict['table_entity']}: {err}"
)
except Exception as exc:
name = value_dict["table_entity"].fullyQualifiedName.__root__
name = value_dict["table_entity"].fullyQualifiedName.root
error = f"Failed to update usage for {name} :{exc}"
logger.debug(traceback.format_exc())
logger.warning(error)
self.status.failed(
StackTraceError(
name=value_dict["table_entity"].fullyQualifiedName.__root__,
name=value_dict["table_entity"].fullyQualifiedName.root,
error=f"Failed to update usage for {name} :{exc}",
stackTrace=traceback.format_exc(),
)
@ -255,7 +255,7 @@ class MetadataUsageBulkSink(BulkSink):
)
)
except Exception as exc:
name = table_entity.name.__root__
name = table_entity.name.root
error = (
f"Error getting usage and join information for {name}: {exc}"
)
@ -281,8 +281,12 @@ class MetadataUsageBulkSink(BulkSink):
"""
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
)
table_joins: TableJoins = TableJoins(
columnJoins=[], directTableJoins=[], startDate=table_usage.date
columnJoins=[], directTableJoins=[], startDate=start_date
)
column_joins_dict = {}
for column_join in table_usage.joins:
@ -317,7 +321,7 @@ class MetadataUsageBulkSink(BulkSink):
key_name = get_column_fqn(table_entity=table_entity, column=key)
if not key_name:
logger.warning(
f"Could not find column {key} in table {table_entity.fullyQualifiedName.__root__}"
f"Could not find column {key} in table {table_entity.fullyQualifiedName.root}"
)
continue
table_joins.columnJoins.append(
@ -370,15 +374,15 @@ class MetadataUsageBulkSink(BulkSink):
query_type = get_query_type(create_query=create_query)
if query_type:
access_details = AccessDetails(
timestamp=create_query.queryDate.__root__,
timestamp=create_query.queryDate.root,
accessedBy=user,
accessedByAProcess=process_user,
)
life_cycle_attr = getattr(life_cycle, query_type)
if (
not life_cycle_attr
or life_cycle_attr.timestamp.__root__
< access_details.timestamp.__root__
or life_cycle_attr.timestamp.root
< access_details.timestamp.root
):
setattr(life_cycle, query_type, access_details)

24
ingestion/src/metadata/ingestion/connections/builders.py
View File

@ -45,8 +45,8 @@ def get_connection_args_common(connection) -> Dict[str, Any]:
"""
return (
connection.connectionArguments.__root__
if connection.connectionArguments and connection.connectionArguments.__root__
connection.connectionArguments.root
if connection.connectionArguments and connection.connectionArguments.root
else {}
)
@ -90,8 +90,8 @@ def get_connection_options_dict(connection) -> Optional[Dict[str, Any]]:
dictionary if exists
"""
return (
connection.connectionOptions.__root__
if connection.connectionOptions and connection.connectionOptions.__root__
connection.connectionOptions.root
if connection.connectionOptions and connection.connectionOptions.root
else None
)
@ -101,12 +101,12 @@ def init_empty_connection_arguments() -> ConnectionArguments:
Initialize a ConnectionArguments model with an empty dictionary.
This helps set keys without further validations.
Running `ConnectionArguments()` returns `ConnectionArguments(__root__=None)`.
Running `ConnectionArguments()` returns `ConnectionArguments(root=None)`.
Instead, we want `ConnectionArguments(__root__={}})` so that
we can pass new keys easily as `connectionArguments.__root__["key"] = "value"`
Instead, we want `ConnectionArguments(root={}})` so that
we can pass new keys easily as `connectionArguments.root["key"] = "value"`
"""
return ConnectionArguments(__root__={})
return ConnectionArguments(root={})
def init_empty_connection_options() -> ConnectionOptions:
@ -114,12 +114,12 @@ def init_empty_connection_options() -> ConnectionOptions:
Initialize a ConnectionOptions model with an empty dictionary.
This helps set keys without further validations.
Running `ConnectionOptions()` returns `ConnectionOptions(__root__=None)`.
Running `ConnectionOptions()` returns `ConnectionOptions(root=None)`.
Instead, we want `ConnectionOptions(__root__={}})` so that
we can pass new keys easily as `ConnectionOptions.__root__["key"] = "value"`
Instead, we want `ConnectionOptions(root={}})` so that
we can pass new keys easily as `ConnectionOptions.root["key"] = "value"`
"""
return ConnectionOptions(__root__={})
return ConnectionOptions(root={})
def _add_password(url: str, connection) -> str:

14
ingestion/src/metadata/ingestion/connections/secrets.py
View File

@ -17,25 +17,27 @@ from functools import wraps
from metadata.ingestion.models.custom_pydantic import CustomSecretStr
# Annotated CustomSecretStr does not like the get_secret_value()
# pylint: disable=no-member
def update_connection_opts_args(connection):
if (
hasattr(connection, "connectionOptions")
and connection.connectionOptions
and connection.connectionOptions.__root__
and connection.connectionOptions.root
):
for key, value in connection.connectionOptions.__root__.items():
for key, value in connection.connectionOptions.root.items():
if isinstance(value, str):
connection.connectionOptions.__root__[key] = CustomSecretStr(
connection.connectionOptions.root[key] = CustomSecretStr(
value
).get_secret_value()
if (
hasattr(connection, "connectionArguments")
and connection.connectionArguments
and connection.connectionArguments.__root__
and connection.connectionArguments.root
):
for key, value in connection.connectionArguments.__root__.items():
for key, value in connection.connectionArguments.root.items():
if isinstance(value, str):
connection.connectionArguments.__root__[key] = CustomSecretStr(
connection.connectionArguments.root[key] = CustomSecretStr(
value
).get_secret_value()

13
ingestion/src/metadata/ingestion/connections/test_connections.py
View File

@ -13,7 +13,7 @@ Classes and methods to handle connection testing when
creating a service
"""
import traceback
from datetime import datetime
from datetime import datetime, timezone
from functools import partial
from typing import Callable, List, Optional
@ -36,6 +36,7 @@ from metadata.generated.schema.entity.services.connections.testConnectionResult
TestConnectionResult,
TestConnectionStepResult,
)
from metadata.generated.schema.type.basic import Timestamp
from metadata.ingestion.ometa.ometa_api import OpenMetadata
from metadata.profiler.orm.functions.conn_test import ConnTestFn
from metadata.utils.logger import cli_logger
@ -146,12 +147,16 @@ def _test_connection_steps_automation_workflow(
# break the workflow if the step is a short circuit step
break
test_connection_result.lastUpdatedAt = datetime.now().timestamp()
test_connection_result.lastUpdatedAt = Timestamp(
int(datetime.now(timezone.utc).timestamp() * 1000)
)
metadata.patch_automation_workflow_response(
automation_workflow, test_connection_result, WorkflowStatus.Running
)
test_connection_result.lastUpdatedAt = datetime.now().timestamp()
test_connection_result.lastUpdatedAt = Timestamp(
int(datetime.now(timezone.utc).timestamp() * 1000)
)
test_connection_result.status = (
StatusType.Failed
@ -169,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().timestamp()
test_connection_result.lastUpdatedAt = datetime.now(tz=timezone.utc).timestamp()
metadata.create_or_update(
CreateWorkflowRequest(
name=automation_workflow.name,

9
ingestion/src/metadata/ingestion/lineage/models.py
View File

@ -14,7 +14,7 @@ Models related to lineage parsing
from enum import Enum
from typing import Dict, List, Optional
from pydantic import BaseModel, Extra, Field
from pydantic import BaseModel, ConfigDict, Field
from metadata.generated.schema.entity.services.connections.database.athenaConnection import (
AthenaType,
@ -144,11 +144,12 @@ class QueryParsingError(BaseModel):
error (str): The error message of the failed query.
"""
class Config:
extra = Extra.forbid
model_config = ConfigDict(
extra="forbid",
)
query: str = Field(..., description="query text of the failed query")
error: Optional[str] = Field(..., description="error message of the failed query")
error: Optional[str] = Field(None, description="error message of the failed query")
class QueryParsingFailures(metaclass=Singleton):

12
ingestion/src/metadata/ingestion/lineage/sql_lineage.py
View File

@ -50,8 +50,8 @@ def get_column_fqn(table_entity: Table, column: str) -> Optional[str]:
if not table_entity:
return None
for tbl_column in table_entity.columns:
if column.lower() == tbl_column.name.__root__.lower():
return tbl_column.fullyQualifiedName.__root__
if column.lower() == tbl_column.name.root.lower():
return tbl_column.fullyQualifiedName.root
return None
@ -226,7 +226,7 @@ def get_column_lineage(
# Select all
if "*" in column_lineage_map.get(to_table_raw_name).get(from_table_raw_name)[0]:
column_lineage_map[to_table_raw_name][from_table_raw_name] = [
(c.name.__root__, c.name.__root__) for c in from_entity.columns
(c.name.root, c.name.root) for c in from_entity.columns
]
# Other cases
@ -268,11 +268,11 @@ def _build_table_lineage(
lineage = AddLineageRequest(
edge=EntitiesEdge(
fromEntity=EntityReference(
id=from_entity.id.__root__,
id=from_entity.id.root,
type="table",
),
toEntity=EntityReference(
id=to_entity.id.__root__,
id=to_entity.id.root,
type="table",
),
)
@ -466,7 +466,7 @@ def get_lineage_via_table_entity(
try:
logger.debug(f"Getting lineage via table entity using query: {query}")
lineage_parser = LineageParser(query, dialect, timeout_seconds=timeout_seconds)
to_table_name = table_entity.name.__root__
to_table_name = table_entity.name.root
for from_table_name in lineage_parser.source_tables:
yield from _create_lineage_by_table_name(

18
ingestion/src/metadata/ingestion/models/custom_properties.py
View File

@ -54,12 +54,12 @@ class CustomPropertyType(BaseModel):
id: basic.Uuid
name: basic.EntityName
displayName: Optional[str]
fullyQualifiedName: Optional[basic.FullyQualifiedEntityName]
description: Optional[basic.Markdown]
category: Optional[str]
nameSpace: Optional[str]
version: Optional[entityHistory.EntityVersion]
updatedAt: Optional[basic.Timestamp]
updatedBy: Optional[str]
href: Optional[basic.Href]
displayName: Optional[str] = None
fullyQualifiedName: Optional[basic.FullyQualifiedEntityName] = None
description: Optional[basic.Markdown] = None
category: Optional[str] = None
nameSpace: Optional[str] = None
version: Optional[entityHistory.EntityVersion] = None
updatedAt: Optional[basic.Timestamp] = None
updatedBy: Optional[str] = None
href: Optional[basic.Href] = None

116
ingestion/src/metadata/ingestion/models/custom_pydantic.py
View File

@ -15,20 +15,71 @@ This classes are used in the generated module, which should have NO
dependencies against any other metadata package. This class should
be self-sufficient with only pydantic at import time.
"""
import json
import logging
import warnings
from typing import Any, Dict
from typing import Any, Dict, Literal, Optional, Union
from pydantic.types import OptionalInt, SecretStr
from pydantic.utils import update_not_none
from pydantic.validators import constr_length_validator, str_validator
from pydantic import BaseModel as PydanticBaseModel
from pydantic import PlainSerializer
from pydantic.main import IncEx
from pydantic.types import SecretStr
from typing_extensions import Annotated
logger = logging.getLogger("metadata")
SECRET = "secret:"
JSON_ENCODERS = "json_encoders"
class CustomSecretStr(SecretStr):
class BaseModel(PydanticBaseModel):
"""
Base model for OpenMetadata generated models.
Specified as `--base-class BASE_CLASS` in the generator.
"""
def model_dump_json( # pylint: disable=too-many-arguments
self,
*,
indent: Optional[int] = None,
include: IncEx = None,
exclude: IncEx = None,
context: Optional[Dict[str, Any]] = None,
by_alias: bool = False,
exclude_unset: bool = True,
exclude_defaults: bool = True,
exclude_none: bool = True,
round_trip: bool = False,
warnings: Union[bool, Literal["none", "warn", "error"]] = True,
serialize_as_any: bool = False,
) -> str:
"""
This is needed due to https://github.com/pydantic/pydantic/issues/8825
We also tried the suggested `serialize` method but it did not
work well with nested models.
This solution is covered in the `test_pydantic_v2` test comparing the
dump results from V1 vs. V2.
"""
return json.dumps(
self.model_dump(
mode="json",
include=include,
exclude=exclude,
context=context,
by_alias=by_alias,
exclude_unset=exclude_unset,
exclude_none=exclude_none,
exclude_defaults=exclude_defaults,
round_trip=round_trip,
warnings=warnings,
serialize_as_any=serialize_as_any,
),
ensure_ascii=True,
)
class _CustomSecretStr(SecretStr):
"""
Custom SecretStr class which use the configured Secrets Manager to retrieve the actual values.
@ -36,48 +87,9 @@ class CustomSecretStr(SecretStr):
in the secrets store.
"""
min_length: OptionalInt = None
max_length: OptionalInt = None
@classmethod
def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None:
update_not_none(
field_schema,
type="string",
writeOnly=True,
format="password",
minLength=cls.min_length,
maxLength=cls.max_length,
)
@classmethod
def __get_validators__(cls) -> "CallableGenerator":
yield cls.validate
yield constr_length_validator
@classmethod
def validate(cls, value: Any) -> "CustomSecretStr":
if isinstance(value, cls):
return value
value = str_validator(value)
return cls(value)
def __init__(self, value: str):
self._secret_value = value
def __repr__(self) -> str:
return f"SecretStr('{self}')"
def __len__(self) -> int:
return len(self._secret_value)
def display(self) -> str:
warnings.warn(
"`secret_str.display()` is deprecated, use `str(secret_str)` instead",
DeprecationWarning,
)
return str(self)
def get_secret_value(self, skip_secret_manager: bool = False) -> str:
"""
This function should only be called after the SecretsManager has properly
@ -108,3 +120,17 @@ class CustomSecretStr(SecretStr):
f"Secret value [{secret_id}] not present in the configured secrets manager: {exc}"
)
return self._secret_value
CustomSecretStr = Annotated[
_CustomSecretStr, PlainSerializer(lambda secret: secret.get_secret_value())
]
def ignore_type_decoder(type_: Any) -> None:
"""Given a type_, add a custom decoder to the BaseModel
to ignore any decoding errors for that type_."""
# We don't import the constants from the constants module to avoid circular imports
BaseModel.model_config[JSON_ENCODERS][type_] = {
lambda v: v.decode("utf-8", "ignore")
}

2
ingestion/src/metadata/ingestion/models/ometa_classification.py
View File

@ -24,6 +24,6 @@ from metadata.generated.schema.type.basic import FullyQualifiedEntityName
class OMetaTagAndClassification(BaseModel):
fqn: Optional[FullyQualifiedEntityName]
fqn: Optional[FullyQualifiedEntityName] = None
classification_request: CreateClassificationRequest
tag_request: CreateTagRequest

12
ingestion/src/metadata/ingestion/models/patch_request.py
View File

@ -36,7 +36,7 @@ class PatchedEntity(BaseModel):
Store the new entity after patch request
"""
new_entity: Optional[Entity]
new_entity: Optional[Entity] = None
ALLOWED_COLUMN_FIELDS = {
@ -343,14 +343,14 @@ def build_patch(
if allowed_fields:
patch = jsonpatch.make_patch(
json.loads(
source.json(
source.model_dump_json(
exclude_unset=True,
exclude_none=True,
include=allowed_fields,
)
),
json.loads(
destination.json(
destination.model_dump_json(
exclude_unset=True,
exclude_none=True,
include=allowed_fields,
@ -359,8 +359,10 @@ def build_patch(
)
else:
patch: jsonpatch.JsonPatch = jsonpatch.make_patch(
json.loads(source.json(exclude_unset=True, exclude_none=True)),
json.loads(destination.json(exclude_unset=True, exclude_none=True)),
json.loads(source.model_dump_json(exclude_unset=True, exclude_none=True)),
json.loads(
destination.model_dump_json(exclude_unset=True, exclude_none=True)
),
)
if not patch:
return None

4
ingestion/src/metadata/ingestion/models/table_metadata.py
View File

@ -26,8 +26,8 @@ class OMetaTableConstraints(BaseModel):
"""
table: Table
foreign_constraints: Optional[List[Dict]]
constraints: Optional[List[TableConstraint]]
foreign_constraints: Optional[List[Dict]] = None
constraints: Optional[List[TableConstraint]] = None
class ColumnTag(BaseModel):

22
ingestion/src/metadata/ingestion/models/topology.py
View File

@ -16,7 +16,7 @@ import threading
from functools import singledispatchmethod
from typing import Any, Dict, Generic, List, Optional, Type, TypeVar
from pydantic import BaseModel, Extra, Field, create_model
from pydantic import BaseModel, ConfigDict, Field, create_model
from metadata.generated.schema.api.data.createStoredProcedure import (
CreateStoredProcedureRequest,
@ -37,8 +37,9 @@ class NodeStage(BaseModel, Generic[T]):
source.
"""
class Config:
extra = Extra.forbid
model_config = ConfigDict(
extra="forbid",
)
# Required fields to define the yielded entity type and the function processing it
type_: Type[T] = Field(
@ -99,8 +100,9 @@ class TopologyNode(BaseModel):
with the updated element from the OM API.
"""
class Config:
extra = Extra.forbid
model_config = ConfigDict(
extra="forbid",
)
producer: str = Field(
...,
@ -128,8 +130,7 @@ class ServiceTopology(BaseModel):
Bounds all service topologies
"""
class Config:
extra = Extra.allow
model_config = ConfigDict(extra="allow")
class TopologyContext(BaseModel):
@ -137,11 +138,10 @@ class TopologyContext(BaseModel):
Bounds all topology contexts
"""
class Config:
extra = Extra.allow
model_config = ConfigDict(extra="allow")
def __repr__(self):
ctx = {key: value.name.__root__ for key, value in self.__dict__.items()}
ctx = {key: value.name.root for key, value in self.__dict__.items()}
return f"TopologyContext({ctx})"
@classmethod
@ -247,7 +247,7 @@ class TopologyContext(BaseModel):
service_name=self.__dict__["database_service"],
database_name=self.__dict__["database"],
schema_name=self.__dict__["database_schema"],
procedure_name=right.name.__root__,
procedure_name=right.name.root,
)

4
ingestion/src/metadata/ingestion/ometa/auth_provider.py
View File

@ -14,7 +14,7 @@ Interface definition for an Auth provider
import os.path
from abc import ABCMeta, abstractmethod
from dataclasses import dataclass
from datetime import datetime
from datetime import datetime, timezone
from dateutil.relativedelta import relativedelta
@ -85,7 +85,7 @@ class OpenMetadataAuthenticationProvider(AuthenticationProvider):
self.config = config
self.security_config: OpenMetadataJWTClientConfig = self.config.securityConfig
self.jwt_token = None
self.expiry = datetime.now() - relativedelta(years=1)
self.expiry = datetime.now(tz=timezone.utc) - relativedelta(years=1)
@classmethod
def create(cls, config: OpenMetadataConnection):

7
ingestion/src/metadata/ingestion/ometa/client_utils.py
View File

@ -18,7 +18,7 @@ from metadata.generated.schema.entity.data.chart import Chart
from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
OpenMetadataConnection,
)
from metadata.generated.schema.type.entityReference import EntityReference
from metadata.generated.schema.type.basic import FullyQualifiedEntityName
from metadata.ingestion.ometa.ometa_api import OpenMetadata
from metadata.utils import fqn
from metadata.utils.logger import ometa_logger
@ -49,7 +49,7 @@ def create_ometa_client(
def get_chart_entities_from_id(
chart_ids: List[str], metadata: OpenMetadata, service_name: str
) -> List[EntityReference]:
) -> List[FullyQualifiedEntityName]:
"""
Method to get the chart entity using get_by_name api
"""
@ -63,6 +63,5 @@ def get_chart_entities_from_id(
),
)
if chart:
entity = EntityReference(id=chart.id, type="chart")
entities.append(entity)
entities.append(chart.fullyQualifiedName)
return entities

6
ingestion/src/metadata/ingestion/ometa/mixins/custom_property_mixin.py
View File

@ -57,7 +57,7 @@ class OMetaCustomPropertyMixin:
resp = self.client.put(
f"/metadata/types/{entity_schema.get('id')}",
data=ometa_custom_property.createCustomPropertyRequest.json(),
data=ometa_custom_property.createCustomPropertyRequest.model_dump_json(),
)
return resp
@ -75,6 +75,4 @@ class OMetaCustomPropertyMixin:
Get the PropertyType for custom properties
"""
custom_property_type = self.get_custom_property_type(data_type=data_type)
return PropertyType(
__root__=EntityReference(id=custom_property_type.id, type="type")
)
return PropertyType(EntityReference(id=custom_property_type.id, type="type"))

4
ingestion/src/metadata/ingestion/ometa/mixins/dashboard_mixin.py
View File

@ -41,7 +41,7 @@ class OMetaDashboardMixin:
:param dashboard_usage_request: Usage data to add
"""
resp = self.client.put(
f"/usage/dashboard/{dashboard.id.__root__}",
data=dashboard_usage_request.json(),
f"/usage/dashboard/{dashboard.id.root}",
data=dashboard_usage_request.model_dump_json(),
)
logger.debug("Published dashboard usage %s", resp)

16
ingestion/src/metadata/ingestion/ometa/mixins/data_insight_mixin.py
View File

@ -44,7 +44,9 @@ class DataInsightMixin:
record (ReportData): report data
"""
resp = self.client.post("/analytics/dataInsights/data", record.json())
resp = self.client.post(
"/analytics/dataInsights/data", record.model_dump_json()
)
return resp
@ -56,7 +58,7 @@ class DataInsightMixin:
record (ReportData): report data
"""
resp = self.client.put(f"/kpi/{fqn}/kpiResult", record.json())
resp = self.client.put(f"/kpi/{fqn}/kpiResult", record.model_dump_json())
return resp
@ -66,7 +68,9 @@ class DataInsightMixin:
) -> List[WebAnalyticEventData]:
"""Get web analytic event"""
resp = self.client.put("/analytics/web/events/collect", event_data.json())
resp = self.client.put(
"/analytics/web/events/collect", event_data.model_dump_json()
)
return resp
@ -127,7 +131,7 @@ class DataInsightMixin:
request_params,
)
return DataInsightChartResult.parse_obj(resp)
return DataInsightChartResult.model_validate(resp)
def get_kpi_result(self, fqn: str, start_ts, end_ts) -> list[KpiResult]:
"""Given FQN return KPI results
@ -146,9 +150,9 @@ class DataInsightMixin:
return [KpiResult(**data) for data in resp["data"]]
def create_kpi(self, create: CreateKpiRequest) -> Kpi:
resp = self.client.post("/kpi", create.json())
resp = self.client.post("/kpi", create.model_dump_json())
return Kpi.parse_obj(resp)
return Kpi.model_validate(resp)
def get_web_analytic_events(
self, event_type: WebAnalyticEventType, start_ts: int, end_ts: int

6
ingestion/src/metadata/ingestion/ometa/mixins/ingestion_pipeline_mixin.py
View File

@ -47,7 +47,7 @@ class OMetaIngestionPipelineMixin:
"""
resp = self.client.put(
f"{self.get_suffix(IngestionPipeline)}/{ingestion_pipeline_fqn}/pipelineStatus",
data=pipeline_status.json(),
data=pipeline_status.model_dump_json(),
)
logger.debug(
f"Created Pipeline Status for pipeline {ingestion_pipeline_fqn}: {resp}"
@ -104,7 +104,9 @@ class OMetaIngestionPipelineMixin:
)
if resp:
return [PipelineStatus.parse_obj(status) for status in resp.get("data")]
return [
PipelineStatus.model_validate(status) for status in resp.get("data")
]
return None
def get_ingestion_pipeline_by_name(

32
ingestion/src/metadata/ingestion/ometa/mixins/lineage_mixin.py
View File

@ -65,11 +65,11 @@ class OMetaLineageMixin(Generic[T]):
def _update_cache(self, request: AddLineageRequest, response: Dict[str, Any]):
try:
for res in response.get("downstreamEdges", []):
if str(request.edge.toEntity.id.__root__) == res.get("toEntity"):
if str(request.edge.toEntity.id.root) == res.get("toEntity"):
search_cache.put(
(
request.edge.fromEntity.id.__root__,
request.edge.toEntity.id.__root__,
request.edge.fromEntity.id.root,
request.edge.toEntity.id.root,
),
{"edge": res.get("lineageDetails")},
)
@ -80,8 +80,8 @@ class OMetaLineageMixin(Generic[T]):
# discard the cache if failed to update
search_cache.put(
(
request.edge.fromEntity.id.__root__,
request.edge.toEntity.id.__root__,
request.edge.fromEntity.id.root,
request.edge.toEntity.id.root,
),
None,
)
@ -96,8 +96,8 @@ class OMetaLineageMixin(Generic[T]):
try:
patch_op_success = False
if check_patch and data.edge.lineageDetails:
from_id = data.edge.fromEntity.id.__root__
to_id = data.edge.toEntity.id.__root__
from_id = data.edge.fromEntity.id.root
to_id = data.edge.toEntity.id.root
edge = self.get_lineage_edge(from_id, to_id)
if edge:
original: AddLineageRequest = deepcopy(data)
@ -131,20 +131,22 @@ class OMetaLineageMixin(Generic[T]):
patch_op_success = True
if patch_op_success is False:
self.client.put(self.get_suffix(AddLineageRequest), data=data.json())
self.client.put(
self.get_suffix(AddLineageRequest), data=data.model_dump_json()
)
except APIError as err:
logger.debug(traceback.format_exc())
logger.error(
"Error %s trying to PUT lineage for %s: %s",
err.status_code,
data.json(),
data.model_dump_json(),
str(err),
)
raise err
from_entity_lineage = self.get_lineage_by_id(
data.edge.fromEntity.type, str(data.edge.fromEntity.id.__root__)
data.edge.fromEntity.type, str(data.edge.fromEntity.id.root)
)
self._update_cache(data, from_entity_lineage)
@ -209,8 +211,8 @@ class OMetaLineageMixin(Generic[T]):
if patch:
self.client.patch(
f"{self.get_suffix(AddLineageRequest)}/{original.edge.fromEntity.type}/"
f"{original.edge.fromEntity.id.__root__}/{original.edge.toEntity.type}"
f"/{original.edge.toEntity.id.__root__}",
f"{original.edge.fromEntity.id.root}/{original.edge.toEntity.type}"
f"/{original.edge.toEntity.id.root}",
data=str(patch),
)
return str(patch)
@ -299,8 +301,8 @@ class OMetaLineageMixin(Generic[T]):
"""
try:
self.client.delete(
f"{self.get_suffix(AddLineageRequest)}/{edge.fromEntity.type}/{edge.fromEntity.id.__root__}/"
f"{edge.toEntity.type}/{edge.toEntity.id.__root__}"
f"{self.get_suffix(AddLineageRequest)}/{edge.fromEntity.type}/{edge.fromEntity.id.root}/"
f"{edge.toEntity.type}/{edge.toEntity.id.root}"
)
except APIError as err:
logger.debug(traceback.format_exc())
@ -328,7 +330,7 @@ class OMetaLineageMixin(Generic[T]):
connection_type = database_service.serviceType.value
add_lineage_request = get_lineage_by_query(
metadata=self,
service_name=database_service.name.__root__,
service_name=database_service.name.root,
dialect=ConnectionTypeDialectMapper.dialect_of(connection_type),
query=sql,
database_name=database_name,

4
ingestion/src/metadata/ingestion/ometa/mixins/mlmodel_mixin.py
View File

@ -90,7 +90,7 @@ class OMetaMlModelMixin(OMetaLineageMixin):
)
)
mlmodel_lineage = self.get_lineage_by_id(MlModel, str(model.id.__root__))
mlmodel_lineage = self.get_lineage_by_id(MlModel, str(model.id.root))
return mlmodel_lineage
@ -151,7 +151,7 @@ class OMetaMlModelMixin(OMetaLineageMixin):
mlHyperParameters=[
MlHyperParameter(
name=key,
value=value,
value=str(value),
)
for key, value in model.get_params().items()
],

33
ingestion/src/metadata/ingestion/ometa/mixins/patch_mixin.py
View File

@ -30,7 +30,7 @@ from metadata.generated.schema.entity.services.connections.testConnectionResult
TestConnectionResult,
)
from metadata.generated.schema.tests.testCase import TestCase, TestCaseParameterValue
from metadata.generated.schema.type.basic import EntityLink
from metadata.generated.schema.type.basic import EntityLink, Markdown
from metadata.generated.schema.type.entityReference import EntityReference
from metadata.generated.schema.type.lifeCycle import LifeCycle
from metadata.generated.schema.type.tagLabel import TagLabel
@ -64,10 +64,7 @@ def update_column_tags(
Inplace update for the incoming column list
"""
for col in columns:
if (
str(col.fullyQualifiedName.__root__).lower()
== column_tag.column_fqn.lower()
):
if str(col.fullyQualifiedName.root).lower() == column_tag.column_fqn.lower():
if operation == PatchOperation.REMOVE:
for tag in col.tags:
if tag.tagFQN == column_tag.tag_label.tagFQN:
@ -92,14 +89,14 @@ def update_column_description(
for col in columns:
# For dbt the column names in OM and dbt are not always in the same case.
# We'll match the column names in case insensitive way
desc_column = col_dict.get(col.fullyQualifiedName.__root__.lower())
desc_column = col_dict.get(col.fullyQualifiedName.root.lower())
if desc_column:
if col.description and not force:
# If the description is already present and force is not passed,
# description will not be overridden
continue
col.description = desc_column.__root__
col.description = desc_column # Keep the Markdown type
if col.children:
update_column_description(col.children, column_descriptions, force)
@ -201,7 +198,7 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
# https://docs.pydantic.dev/latest/usage/exporting_models/#modelcopy
destination = source.copy(deep=True)
destination.description = description
destination.description = Markdown(description)
return self.patch(entity=entity, source=source, destination=destination)
@ -256,7 +253,7 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
destination = source.copy(deep=True)
destination.entityLink = EntityLink(__root__=entity_link)
destination.entityLink = EntityLink(entity_link)
if test_case_parameter_values:
destination.parameterValues = test_case_parameter_values
if compute_passed_failed_row_count != source.computePassedFailedRowCount:
@ -294,11 +291,11 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
source.tags = instance.tags or []
destination = source.copy(deep=True)
tag_fqns = {label.tagFQN.__root__ for label in tag_labels}
tag_fqns = {label.tagFQN.root for label in tag_labels}
if operation == PatchOperation.REMOVE:
for tag in destination.tags:
if tag.tagFQN.__root__ in tag_fqns:
if tag.tagFQN.root in tag_fqns:
destination.tags.remove(tag)
else:
destination.tags.extend(tag_labels)
@ -394,7 +391,7 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
if patched_entity is None:
logger.debug(
f"Empty PATCH result. Either everything is up to date or the "
f"column names are not in [{table.fullyQualifiedName.__root__}]"
f"column names are not in [{table.fullyQualifiedName.root}]"
)
return patched_entity
@ -440,7 +437,9 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
return self.patch_column_descriptions(
table=table,
column_descriptions=[
ColumnDescription(column_fqn=column_fqn, description=description)
ColumnDescription(
column_fqn=column_fqn, description=Markdown(description)
)
],
force=force,
)
@ -478,7 +477,7 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
if patched_entity is None:
logger.debug(
f"Empty PATCH result. Either everything is up to date or "
f"columns are not matching for [{table.fullyQualifiedName.__root__}]"
f"columns are not matching for [{table.fullyQualifiedName.root}]"
)
return patched_entity
@ -494,7 +493,7 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
"""
result_data: Dict = {
PatchField.PATH: PatchPath.RESPONSE,
PatchField.VALUE: test_connection_result.dict(),
PatchField.VALUE: test_connection_result.model_dump(),
PatchField.OPERATION: PatchOperation.ADD,
}
@ -538,7 +537,7 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
except Exception as exc:
logger.debug(traceback.format_exc())
logger.warning(
f"Error trying to Patch life cycle data for {entity.fullyQualifiedName.__root__}: {exc}"
f"Error trying to Patch life cycle data for {entity.fullyQualifiedName.root}: {exc}"
)
return None
@ -553,6 +552,6 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
except Exception as exc:
logger.debug(traceback.format_exc())
logger.warning(
f"Error trying to Patch Domain for {entity.fullyQualifiedName.__root__}: {exc}"
f"Error trying to Patch Domain for {entity.fullyQualifiedName.root}: {exc}"
)
return None

2
ingestion/src/metadata/ingestion/ometa/mixins/pipeline_mixin.py
View File

@ -43,7 +43,7 @@ class OMetaPipelineMixin:
"""
resp = self.client.put(
f"{self.get_suffix(Pipeline)}/{fqn}/status",
data=status.json(),
data=status.model_dump_json(),
)
return Pipeline(**resp)

8
ingestion/src/metadata/ingestion/ometa/mixins/query_mixin.py
View File

@ -41,10 +41,10 @@ class OMetaQueryMixin:
return str(result.hexdigest())
def _get_or_create_query(self, query: CreateQueryRequest) -> Optional[Query]:
query_hash = self._get_query_hash(query=query.query.__root__)
query_hash = self._get_query_hash(query=query.query.root)
query_entity = self.get_by_name(entity=Query, fqn=query_hash)
if query_entity is None:
resp = self.client.put(self.get_suffix(Query), data=query.json())
resp = self.client.put(self.get_suffix(Query), data=query.model_dump_json())
if resp and resp.get("id"):
query_entity = Query(**resp)
return query_entity
@ -63,9 +63,9 @@ class OMetaQueryMixin:
query = self._get_or_create_query(create_query)
if query:
# Add Query Usage
table_ref = EntityReference(id=entity.id.__root__, type="table")
table_ref = EntityReference(id=entity.id.root, type="table")
# convert object to json array string
table_ref_json = "[" + table_ref.json() + "]"
table_ref_json = "[" + table_ref.model_dump_json() + "]"
self.client.put(
f"{self.get_suffix(Query)}/{model_str(query.id)}/usage",
data=table_ref_json,

26
ingestion/src/metadata/ingestion/ometa/mixins/role_policy_mixin.py
View File

@ -125,7 +125,7 @@ class OMetaRolePolicyMixin(OMetaPatchMixinBase):
if previous is None
else PatchOperation.REPLACE,
PatchField.PATH: path.format(rule_index=rule_index),
PatchField.VALUE: str(current.__root__),
PatchField.VALUE: str(current.root),
}
]
return data
@ -158,10 +158,10 @@ class OMetaRolePolicyMixin(OMetaPatchMixinBase):
if not instance:
return None
policy_index: int = len(instance.policies.__root__) - 1
policy_index: int = len(instance.policies.root) - 1
data: List
if operation is PatchOperation.REMOVE:
if len(instance.policies.__root__) == 1:
if len(instance.policies.root) == 1:
logger.error(
f"The Role with id [{model_str(entity_id)}] has only one (1)"
f" policy. Unable to remove."
@ -177,7 +177,7 @@ class OMetaRolePolicyMixin(OMetaPatchMixinBase):
index: int = 0
is_policy_found: bool = False
for policy in instance.policies.__root__:
for policy in instance.policies.root:
if model_str(policy.id) == model_str(policy_id):
is_policy_found = True
continue
@ -187,7 +187,7 @@ class OMetaRolePolicyMixin(OMetaPatchMixinBase):
PatchField.PATH: PatchPath.POLICIES_DESCRIPTION.format(
index=index
),
PatchField.VALUE: model_str(policy.description.__root__),
PatchField.VALUE: model_str(policy.description.root),
}
)
data.append(
@ -294,7 +294,7 @@ class OMetaRolePolicyMixin(OMetaPatchMixinBase):
if not instance:
return None
rule_index: int = len(instance.rules.__root__) - 1
rule_index: int = len(instance.rules.root) - 1
data: List[Dict]
if operation == PatchOperation.ADD:
data = [
@ -303,7 +303,7 @@ class OMetaRolePolicyMixin(OMetaPatchMixinBase):
PatchField.PATH: PatchPath.RULES.format(rule_index=rule_index + 1),
PatchField.VALUE: {
PatchValue.NAME: rule.name,
PatchValue.CONDITION: rule.condition.__root__,
PatchValue.CONDITION: rule.condition.root,
PatchValue.EFFECT: rule.effect.name,
PatchValue.OPERATIONS: [
operation.name for operation in rule.operations
@ -314,12 +314,12 @@ class OMetaRolePolicyMixin(OMetaPatchMixinBase):
]
if rule.description is not None:
data[0][PatchField.VALUE][PatchValue.DESCRIPTION] = str(
rule.description.__root__
rule.description.root
)
if rule.fullyQualifiedName is not None:
data[0][PatchField.VALUE][PatchValue.FQN] = str(
rule.fullyQualifiedName.__root__
rule.fullyQualifiedName.root
)
else:
@ -334,8 +334,8 @@ class OMetaRolePolicyMixin(OMetaPatchMixinBase):
}
]
for rule_index in range(len(instance.rules.__root__) - 1, -1, -1):
current_rule: Rule = instance.rules.__root__[rule_index]
for rule_index in range(len(instance.rules.root) - 1, -1, -1):
current_rule: Rule = instance.rules.root[rule_index]
if current_rule.name == rule.name:
break
@ -345,7 +345,7 @@ class OMetaRolePolicyMixin(OMetaPatchMixinBase):
)
return None
previous_rule: Rule = instance.rules.__root__[rule_index - 1]
previous_rule: Rule = instance.rules.root[rule_index - 1]
# Condition
data.append(
{
@ -353,7 +353,7 @@ class OMetaRolePolicyMixin(OMetaPatchMixinBase):
PatchField.PATH: PatchPath.RULES_CONDITION.format(
rule_index=rule_index - 1
),
PatchField.VALUE: current_rule.condition.__root__,
PatchField.VALUE: current_rule.condition.root,
}
)
# Description - Optional

10
ingestion/src/metadata/ingestion/ometa/mixins/search_index_mixin.py
View File

@ -47,13 +47,13 @@ class OMetaSearchIndexMixin:
resp = None
try:
resp = self.client.put(
f"{self.get_suffix(SearchIndex)}/{search_index.id.__root__}/sampleData",
data=sample_data.json(),
f"{self.get_suffix(SearchIndex)}/{search_index.id.root}/sampleData",
data=sample_data.model_dump_json(),
)
except Exception as exc:
logger.debug(traceback.format_exc())
logger.warning(
f"Error trying to PUT sample data for {search_index.fullyQualifiedName.__root__}: {exc}"
f"Error trying to PUT sample data for {search_index.fullyQualifiedName.root}: {exc}"
)
if resp:
@ -63,13 +63,13 @@ class OMetaSearchIndexMixin:
logger.debug(traceback.format_exc())
logger.warning(
"Unicode Error parsing the sample data response "
f"from {search_index.fullyQualifiedName.__root__}: {err}"
f"from {search_index.fullyQualifiedName.root}: {err}"
)
except Exception as exc:
logger.debug(traceback.format_exc())
logger.warning(
"Error trying to parse sample data results"
f"from {search_index.fullyQualifiedName.__root__}: {exc}"
f"from {search_index.fullyQualifiedName.root}: {exc}"
)
return None

9
ingestion/src/metadata/ingestion/ometa/mixins/server_mixin.py
View File

@ -21,7 +21,6 @@ from metadata.__version__ import (
match_versions,
)
from metadata.generated.schema.settings.settings import Settings, SettingType
from metadata.ingestion.models.encoders import show_secrets_encoder
from metadata.ingestion.ometa.client import REST
from metadata.ingestion.ometa.routes import ROUTES
from metadata.utils.logger import ometa_logger
@ -91,9 +90,9 @@ class OMetaServerMixin:
Returns:
Settings
"""
data = settings.json(encoder=show_secrets_encoder)
data = settings.model_dump_json()
response = self.client.put(ROUTES.get(Settings.__name__), data)
return Settings.parse_obj(response)
return Settings.model_validate(response)
def get_settings_by_name(self, setting_type: SettingType) -> Optional[Settings]:
"""Get setting by name
@ -106,7 +105,7 @@ class OMetaServerMixin:
)
if not response:
return None
return Settings.parse_obj(response)
return Settings.model_validate(response)
def get_profiler_config_settings(self) -> Optional[Settings]:
"""Get profiler config setting
@ -117,4 +116,4 @@ class OMetaServerMixin:
response = self.client.get("/system/settings/profilerConfiguration")
if not response:
return None
return Settings.parse_obj(response)
return Settings.model_validate(response)

4
ingestion/src/metadata/ingestion/ometa/mixins/service_mixin.py
View File

@ -55,8 +55,8 @@ class OMetaServiceMixin:
create_entity_class = self.get_create_entity_type(entity=entity)
return create_entity_class(
name=config.serviceName,
serviceType=config.serviceConnection.__root__.config.type.value,
connection=config.serviceConnection.__root__
serviceType=config.serviceConnection.root.config.type.value,
connection=config.serviceConnection.root
if self.config.storeServiceConnection
else None,
)

4
ingestion/src/metadata/ingestion/ometa/mixins/suggestions_mixin.py
View File

@ -34,8 +34,8 @@ class OMetaSuggestionsMixin:
Update an existing Suggestion with new fields
"""
resp = self.client.put(
f"{self.get_suffix(Suggestion)}/{str(suggestion.id.__root__)}",
data=suggestion.json(),
f"{self.get_suffix(Suggestion)}/{str(suggestion.root.id.root)}",
data=suggestion.model_dump_json(),
)
return Suggestion(**resp)

Some files were not shown because too many files have changed in this diff Show More