acryl-datahub-gx-plugin 1.3.1.5rc10__py3-none-any.whl → 1.3.1.6__py3-none-any.whl

{acryl_datahub_gx_plugin-1.3.1.5rc10.dist-info → acryl_datahub_gx_plugin-1.3.1.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: acryl-datahub-gx-plugin
-Version: 1.3.1.5rc10
+Version: 1.3.1.6
 Summary: Datahub GX plugin to capture executions and send to Datahub
 Home-page: https://docs.datahub.com/
 License: Apache-2.0
@@ -22,87 +22,87 @@ Classifier: Topic :: Software Development
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 Requires-Dist: great-expectations<1.0.0,>=0.17.15
+Requires-Dist: pydantic>=2.1.0
+Requires-Dist: requests
+Requires-Dist: sqlalchemy<2,>=1.4.39
+Requires-Dist: acryl-datahub[datahub-rest,sql-parser]==1.3.1.6
 Requires-Dist: traitlets!=5.2.2
 Requires-Dist: requests_file
-Requires-Dist: sqlalchemy<2,>=1.4.39
-Requires-Dist: requests
-Requires-Dist: pydantic>=2.1.0
-Requires-Dist: acryl-datahub[datahub-rest,sql-parser]==1.3.1.5rc10
 Provides-Extra: ignore
 Provides-Extra: dev
-Requires-Dist: requests_file; extra == "dev"
-Requires-Dist: jsonpickle; extra == "dev"
+Requires-Dist: mypy==1.17.1; extra == "dev"
+Requires-Dist: types-freezegun; extra == "dev"
+Requires-Dist: acryl-datahub[datahub-rest,sql-parser]==1.3.1.6; extra == "dev"
+Requires-Dist: build; extra == "dev"
+Requires-Dist: twine; extra == "dev"
+Requires-Dist: types-setuptools; extra == "dev"
+Requires-Dist: types-click==0.1.12; extra == "dev"
+Requires-Dist: types-toml; extra == "dev"
+Requires-Dist: types-cachetools; extra == "dev"
+Requires-Dist: types-tabulate; extra == "dev"
 Requires-Dist: pydantic>=2.1.0; extra == "dev"
-Requires-Dist: tox; extra == "dev"
-Requires-Dist: deepdiff!=8.0.0; extra == "dev"
-Requires-Dist: ruff==0.11.7; extra == "dev"
 Requires-Dist: pytest-asyncio>=0.16.0; extra == "dev"
-Requires-Dist: types-PyYAML; extra == "dev"
+Requires-Dist: types-six; extra == "dev"
+Requires-Dist: tox; extra == "dev"
+Requires-Dist: packaging; extra == "dev"
+Requires-Dist: freezegun; extra == "dev"
 Requires-Dist: types-requests; extra == "dev"
-Requires-Dist: types-cachetools; extra == "dev"
+Requires-Dist: ruff==0.11.7; extra == "dev"
+Requires-Dist: types-pytz; extra == "dev"
+Requires-Dist: types-dataclasses; extra == "dev"
+Requires-Dist: coverage>=5.1; extra == "dev"
+Requires-Dist: pytest>=6.2.2; extra == "dev"
+Requires-Dist: sqlalchemy<2,>=1.4.39; extra == "dev"
 Requires-Dist: sqlalchemy-stubs; extra == "dev"
 Requires-Dist: types-python-dateutil; extra == "dev"
-Requires-Dist: types-dataclasses; extra == "dev"
-Requires-Dist: packaging; extra == "dev"
-Requires-Dist: mypy==1.17.1; extra == "dev"
-Requires-Dist: types-freezegun; extra == "dev"
-Requires-Dist: requests; extra == "dev"
-Requires-Dist: types-pytz; extra == "dev"
-Requires-Dist: pytest-cov>=2.8.1; extra == "dev"
 Requires-Dist: traitlets!=5.2.2; extra == "dev"
-Requires-Dist: types-six; extra == "dev"
-Requires-Dist: types-setuptools; extra == "dev"
+Requires-Dist: requests_file; extra == "dev"
+Requires-Dist: pytest-cov>=2.8.1; extra == "dev"
+Requires-Dist: deepdiff!=8.0.0; extra == "dev"
+Requires-Dist: requests; extra == "dev"
 Requires-Dist: great-expectations<1.0.0,>=0.17.15; extra == "dev"
-Requires-Dist: types-tabulate; extra == "dev"
-Requires-Dist: pytest>=6.2.2; extra == "dev"
-Requires-Dist: types-toml; extra == "dev"
-Requires-Dist: acryl-datahub[datahub-rest,sql-parser]==1.3.1.5rc10; extra == "dev"
-Requires-Dist: twine; extra == "dev"
-Requires-Dist: freezegun; extra == "dev"
+Requires-Dist: types-PyYAML; extra == "dev"
+Requires-Dist: jsonpickle; extra == "dev"
 Requires-Dist: requests-mock; extra == "dev"
-Requires-Dist: sqlalchemy<2,>=1.4.39; extra == "dev"
-Requires-Dist: coverage>=5.1; extra == "dev"
-Requires-Dist: build; extra == "dev"
-Requires-Dist: types-click==0.1.12; extra == "dev"
 Provides-Extra: integration-tests
-Requires-Dist: acryl-datahub[testing-utils]==1.3.1.5rc10; extra == "integration-tests"
-Requires-Dist: requests_file; extra == "integration-tests"
-Requires-Dist: pytest-docker>=1.1.0; extra == "integration-tests"
-Requires-Dist: jsonpickle; extra == "integration-tests"
+Requires-Dist: mypy==1.17.1; extra == "integration-tests"
+Requires-Dist: pyspark; extra == "integration-tests"
+Requires-Dist: types-freezegun; extra == "integration-tests"
+Requires-Dist: acryl-datahub[datahub-rest,sql-parser]==1.3.1.6; extra == "integration-tests"
+Requires-Dist: build; extra == "integration-tests"
+Requires-Dist: twine; extra == "integration-tests"
+Requires-Dist: types-setuptools; extra == "integration-tests"
+Requires-Dist: types-click==0.1.12; extra == "integration-tests"
+Requires-Dist: types-toml; extra == "integration-tests"
+Requires-Dist: types-cachetools; extra == "integration-tests"
+Requires-Dist: types-tabulate; extra == "integration-tests"
 Requires-Dist: pydantic>=2.1.0; extra == "integration-tests"
-Requires-Dist: tox; extra == "integration-tests"
-Requires-Dist: deepdiff!=8.0.0; extra == "integration-tests"
-Requires-Dist: ruff==0.11.7; extra == "integration-tests"
 Requires-Dist: pytest-asyncio>=0.16.0; extra == "integration-tests"
-Requires-Dist: types-PyYAML; extra == "integration-tests"
+Requires-Dist: types-six; extra == "integration-tests"
+Requires-Dist: tox; extra == "integration-tests"
+Requires-Dist: packaging; extra == "integration-tests"
+Requires-Dist: freezegun; extra == "integration-tests"
 Requires-Dist: types-requests; extra == "integration-tests"
-Requires-Dist: psycopg2-binary; extra == "integration-tests"
-Requires-Dist: types-cachetools; extra == "integration-tests"
+Requires-Dist: ruff==0.11.7; extra == "integration-tests"
+Requires-Dist: types-pytz; extra == "integration-tests"
+Requires-Dist: types-dataclasses; extra == "integration-tests"
+Requires-Dist: coverage>=5.1; extra == "integration-tests"
+Requires-Dist: pytest>=6.2.2; extra == "integration-tests"
+Requires-Dist: sqlalchemy<2,>=1.4.39; extra == "integration-tests"
 Requires-Dist: sqlalchemy-stubs; extra == "integration-tests"
 Requires-Dist: types-python-dateutil; extra == "integration-tests"
-Requires-Dist: types-dataclasses; extra == "integration-tests"
-Requires-Dist: packaging; extra == "integration-tests"
-Requires-Dist: mypy==1.17.1; extra == "integration-tests"
-Requires-Dist: types-freezegun; extra == "integration-tests"
-Requires-Dist: requests; extra == "integration-tests"
-Requires-Dist: types-pytz; extra == "integration-tests"
-Requires-Dist: pytest-cov>=2.8.1; extra == "integration-tests"
 Requires-Dist: traitlets!=5.2.2; extra == "integration-tests"
-Requires-Dist: types-six; extra == "integration-tests"
-Requires-Dist: types-setuptools; extra == "integration-tests"
+Requires-Dist: requests_file; extra == "integration-tests"
+Requires-Dist: pytest-cov>=2.8.1; extra == "integration-tests"
+Requires-Dist: deepdiff!=8.0.0; extra == "integration-tests"
+Requires-Dist: psycopg2-binary; extra == "integration-tests"
+Requires-Dist: requests; extra == "integration-tests"
+Requires-Dist: acryl-datahub[testing-utils]==1.3.1.6; extra == "integration-tests"
+Requires-Dist: pytest-docker>=1.1.0; extra == "integration-tests"
 Requires-Dist: great-expectations<1.0.0,>=0.17.15; extra == "integration-tests"
-Requires-Dist: types-tabulate; extra == "integration-tests"
-Requires-Dist: pytest>=6.2.2; extra == "integration-tests"
-Requires-Dist: types-toml; extra == "integration-tests"
-Requires-Dist: pyspark; extra == "integration-tests"
-Requires-Dist: acryl-datahub[datahub-rest,sql-parser]==1.3.1.5rc10; extra == "integration-tests"
-Requires-Dist: twine; extra == "integration-tests"
-Requires-Dist: freezegun; extra == "integration-tests"
+Requires-Dist: types-PyYAML; extra == "integration-tests"
+Requires-Dist: jsonpickle; extra == "integration-tests"
 Requires-Dist: requests-mock; extra == "integration-tests"
-Requires-Dist: sqlalchemy<2,>=1.4.39; extra == "integration-tests"
-Requires-Dist: coverage>=5.1; extra == "integration-tests"
-Requires-Dist: build; extra == "integration-tests"
-Requires-Dist: types-click==0.1.12; extra == "integration-tests"
 Dynamic: classifier
 Dynamic: description
 Dynamic: description-content-type

acryl_datahub_gx_plugin-1.3.1.6.dist-info/RECORD

@@ -0,0 +1,8 @@
+datahub_gx_plugin/__init__.py,sha256=EAwU2Z4u9abD5xHwcfgwo1xvK8-lSfByM-cp0IMw4mQ,69
+datahub_gx_plugin/_version.py,sha256=GlHJansOn_-rYbo_oQsFvc5W6jU54Ppp5V4ORtwd-FY,135
+datahub_gx_plugin/action.py,sha256=d6H0daSN0SrZ46ZfNPdVwldQYk1eJcSryHSC6cWnjM8,40665
+acryl_datahub_gx_plugin-1.3.1.6.dist-info/METADATA,sha256=RK93bGb4Ku1gqlyLlVX6j1Qn5nkpmLHwF4ceq-sKOas,5639
+acryl_datahub_gx_plugin-1.3.1.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+acryl_datahub_gx_plugin-1.3.1.6.dist-info/entry_points.txt,sha256=MPipKHfWSerlcRwWt8OP93TPCVxpposDMyILo9cszbM,88
+acryl_datahub_gx_plugin-1.3.1.6.dist-info/top_level.txt,sha256=ryn3FMtO1isrM-TbO-SRYy_-31fUpgdEBmBSrAmvL1c,18
+acryl_datahub_gx_plugin-1.3.1.6.dist-info/RECORD,,

datahub_gx_plugin/_version.py

@@ -1,3 +1,3 @@
 # Published at https://pypi.org/project/acryl-datahub-gx-plugin/.
 __package_name__ = "acryl-datahub-gx-plugin"
-__version__ = "1.3.1.5rc10"
+__version__ = "1.3.1.6"

datahub_gx_plugin/action.py

@@ -25,7 +25,10 @@ from great_expectations.data_context.types.resource_identifiers import (
     ExpectationSuiteIdentifier,
     ValidationResultIdentifier,
 )
-from great_expectations.execution_engine import PandasExecutionEngine
+from great_expectations.execution_engine import (
+    PandasExecutionEngine,
+    SparkDFExecutionEngine,
+)
 from great_expectations.execution_engine.sqlalchemy_execution_engine import (
     SqlAlchemyExecutionEngine,
 )
@@ -586,16 +589,67 @@ class DataHubValidationAction(ValidationAction):
         )
 
     def get_dataset_partitions(self, batch_identifier, data_asset):
-        dataset_partitions = []
+        dataset_partitions: List[
+            Dict[str, Union[PartitionSpecClass, BatchSpec, str, None]]
+        ] = []
 
         logger.debug("Finding datasets being validated")
 
-        # for now, we support only v3-api and sqlalchemy execution engine and Pandas engine
+        # for now, we support only v3-api and sqlalchemy execution engine,Pandas engine and Spark engine
         is_sql_alchemy = isinstance(data_asset, Validator) and (
             isinstance(data_asset.execution_engine, SqlAlchemyExecutionEngine)
         )
         is_pandas = isinstance(data_asset.execution_engine, PandasExecutionEngine)
-        if is_sql_alchemy or is_pandas:
+
+        is_spark = isinstance(data_asset.execution_engine, SparkDFExecutionEngine)
+
+        if is_spark:
+            ge_batch_spec = data_asset.active_batch_spec
+            partitionSpec = None
+            batchSpecProperties = {
+                "data_asset_name": str(
+                    data_asset.active_batch_definition.data_asset_name
+                ),
+                "datasource_name": str(
+                    data_asset.active_batch_definition.datasource_name
+                ),
+            }
+
+            if isinstance(ge_batch_spec, RuntimeDataBatchSpec):
+                data_platform = self.get_platform_instance_spark(
+                    data_asset.active_batch_definition.datasource_name
+                )
+
+                dataset_urn = builder.make_dataset_urn_with_platform_instance(
+                    platform=(
+                        data_platform
+                        if self.platform_alias is None
+                        else self.platform_alias
+                    ),
+                    name=data_asset.active_batch_definition.data_asset_name,
+                    platform_instance="",
+                    env=self.env,
+                )
+
+                batchSpec = BatchSpec(
+                    nativeBatchId=batch_identifier,
+                    query="",
+                    customProperties=batchSpecProperties,
+                )
+                dataset_partitions.append(
+                    {
+                        "dataset_urn": dataset_urn,
+                        "partitionSpec": partitionSpec,
+                        "batchSpec": batchSpec,
+                    }
+                )
+            else:
+                warn(
+                    "DataHubValidationAction does not recognize this GE batch spec type for SparkDFExecutionEngine- {batch_spec_type}. No action will be taken.".format(
+                        batch_spec_type=type(ge_batch_spec)
+                    )
+                )
+        elif is_sql_alchemy or is_pandas:
             ge_batch_spec = data_asset.active_batch_spec
             partitionSpec = None
             batchSpecProperties = {
@@ -607,6 +661,7 @@ class DataHubValidationAction(ValidationAction):
                 ),
             }
             sqlalchemy_uri = None
+
             if is_sql_alchemy and isinstance(
                 data_asset.execution_engine.engine, Engine
             ):
@@ -627,7 +682,7 @@ class DataHubValidationAction(ValidationAction):
                     schema_name,
                     table_name,
                     self.env,
-                    self.get_platform_instance(
+                    self.get_platform_instance_sqlalchemy(
                         data_asset.active_batch_definition.datasource_name
                     ),
                     self.exclude_dbname,
@@ -709,7 +764,7 @@ class DataHubValidationAction(ValidationAction):
                         None,
                         table,
                         self.env,
-                        self.get_platform_instance(
+                        self.get_platform_instance_sqlalchemy(
                             data_asset.active_batch_definition.datasource_name
                         ),
                         self.exclude_dbname,
@@ -724,7 +779,7 @@ class DataHubValidationAction(ValidationAction):
                         }
                     )
             elif isinstance(ge_batch_spec, RuntimeDataBatchSpec):
-                data_platform = self.get_platform_instance(
+                data_platform = self.get_platform_instance_sqlalchemy(
                     data_asset.active_batch_definition.datasource_name
                 )
                 dataset_urn = builder.make_dataset_urn_with_platform_instance(
@@ -758,14 +813,14 @@ class DataHubValidationAction(ValidationAction):
         else:
             # TODO - v2-spec - SqlAlchemyDataset support
             warn(
-                "DataHubValidationAction does not recognize this GE data asset type - {asset_type}. This is either using v2-api or execution engine other than sqlalchemy.".format(
+                "DataHubValidationAction does not recognize this GE data asset type - {asset_type}.".format(
                     asset_type=type(data_asset)
                 )
             )
 
         return dataset_partitions
 
-    def get_platform_instance(self, datasource_name):
+    def get_platform_instance_sqlalchemy(self, datasource_name):
         if self.platform_instance_map and datasource_name in self.platform_instance_map:
             return self.platform_instance_map[datasource_name]
         else:
@@ -774,6 +829,16 @@ class DataHubValidationAction(ValidationAction):
             )
         return None
 
+    def get_platform_instance_spark(self, datasource_name):
+        if self.platform_instance_map and datasource_name in self.platform_instance_map:
+            return self.platform_instance_map[datasource_name]
+        else:
+            warn(
+                f"Datasource {datasource_name} is not present in platform_instance_map. \
+                        Data platform will be {datasource_name} by default "
+            )
+            return datasource_name
+
 
 def parse_int_or_default(value, default_value=None):
     if value is None:

acryl_datahub_gx_plugin-1.3.1.5rc10.dist-info/RECORD

@@ -1,8 +0,0 @@
-datahub_gx_plugin/__init__.py,sha256=EAwU2Z4u9abD5xHwcfgwo1xvK8-lSfByM-cp0IMw4mQ,69
-datahub_gx_plugin/_version.py,sha256=SO5hwP7jAVhc9uGQ1pfby0NKv0UnFGAtY_XjJAVCn3A,139
-datahub_gx_plugin/action.py,sha256=W1TZEQzURdg-QOEUdayQmoGB2oO54UkRs9PQEbxhp7A,38149
-acryl_datahub_gx_plugin-1.3.1.5rc10.dist-info/METADATA,sha256=yNd7izc5_4qKF9qoUZvSvlrHD6C4UrstcNSq61kY5cg,5659
-acryl_datahub_gx_plugin-1.3.1.5rc10.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-acryl_datahub_gx_plugin-1.3.1.5rc10.dist-info/entry_points.txt,sha256=MPipKHfWSerlcRwWt8OP93TPCVxpposDMyILo9cszbM,88
-acryl_datahub_gx_plugin-1.3.1.5rc10.dist-info/top_level.txt,sha256=ryn3FMtO1isrM-TbO-SRYy_-31fUpgdEBmBSrAmvL1c,18
-acryl_datahub_gx_plugin-1.3.1.5rc10.dist-info/RECORD,,