dcs-sdk 1.6.5__tar.gz → 1.6.7__tar.gz

{dcs_sdk-1.6.5 → dcs_sdk-1.6.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dcs-sdk
-Version: 1.6.5
+Version: 1.6.7
 Summary: SDK for DataChecks
 Author: Waterdip Labs
 Author-email: hello@waterdip.ai
@@ -30,6 +30,8 @@ Provides-Extra: sybase
 Provides-Extra: trino
 Provides-Extra: vertica
 Requires-Dist: attrs (>=23.1.0)
+Requires-Dist: azure-identity (>=1.25.1,<2.0.0)
+Requires-Dist: azure-storage-blob (>=12.27.1,<13.0.0)
 Requires-Dist: click (>=8.1)
 Requires-Dist: clickhouse-driver (>=0.2.9) ; extra == "clickhouse" or extra == "all-dbs"
 Requires-Dist: cryptography (>=44.0.1) ; extra == "snowflake" or extra == "all-dbs"
@@ -84,7 +86,7 @@ Requires-Dist: vertica-python (>=1.4.0) ; extra == "vertica" or extra == "all-db
 Description-Content-Type: text/markdown
 
 <h1 align="center">
-  DCS SDK v1.6.4
+  DCS SDK v1.6.7
 </h1>
 
 > SDK for DataChecks

{dcs_sdk-1.6.5 → dcs_sdk-1.6.7}/README.md

@@ -1,5 +1,5 @@
 <h1 align="center">
-  DCS SDK v1.6.4
+  DCS SDK v1.6.7
 </h1>
 
 > SDK for DataChecks

{dcs_sdk-1.6.5 → dcs_sdk-1.6.7}/data_diff/__init__.py

@@ -55,9 +55,7 @@ def connect_to_table(
             db_info.pop(k)
     if isinstance(key_columns, str):
         key_columns = (key_columns,)
-
     db: Database = connect(db_info, thread_count=thread_count)
-
     if isinstance(table_name, str):
         table_name = db.dialect.parse_table_name(table_name)
 

{dcs_sdk-1.6.5 → dcs_sdk-1.6.7}/dcs_core/core/common/errors.py

@@ -16,6 +16,8 @@ ERROR_RUNTIME = "runtime_error"
 ERROR_CONFIGURATION = "configuration_error"
 ERROR_DATA_SOURCES_CONNECTION = "data_sources_connection_error"
 ERROR_METRIC_GENERATION = "metric_generation_error"
+ERROR_FETCHING_TABLE = "table_fetch_error"
+ERROR_FETCHING_COLUMN = "column_fetch_error"
 
 
 class DataChecksRuntimeError(Exception):
@@ -48,3 +50,19 @@ class DataChecksMetricGenerationError(Exception):
     def __init__(self, message):
         super().__init__(message)
         self.error_code = ERROR_METRIC_GENERATION
+
+
+class DatachecksTableFetchError(Exception):
+    """Raised when there is an error in fetching table."""
+
+    def __init__(self, message):
+        super().__init__(message)
+        self.error_code = ERROR_FETCHING_TABLE
+
+
+class DatachecksColumnFetchError(Exception):
+    """Raised when there is an error in fetching column."""
+
+    def __init__(self, message):
+        super().__init__(message)
+        self.error_code = ERROR_FETCHING_COLUMN

{dcs_sdk-1.6.5 → dcs_sdk-1.6.7}/dcs_core/core/common/models/configuration.py

@@ -43,6 +43,7 @@ class DataSourceType(str, Enum):
     ORACLE = "oracle"
     DB2 = "db2"
     SYBASE = "sybase"
+    AZURE_BLOB = "azure_blob"
 
 
 class DataSourceLanguageSupport(str, Enum):
@@ -85,6 +86,11 @@ class DataSourceConnectionConfiguration:
     security: Optional[str] = None  # IBM DB2 specific configuration
     protocol: Optional[str] = None  # IBM DB2 specific configuration
     server: Optional[str] = None
+    account_name: Optional[str] = None
+    container_name: Optional[str] = None
+    account_key: Optional[str] = None
+    endpoint_suffix: Optional[str] = None
+    subfolder_path: Optional[str] = None
 
 
 @dataclass

dcs_sdk-1.6.7/dcs_core/core/datasource/file_datasource.py

@@ -0,0 +1,26 @@
+#  Copyright 2022-present, the Waterdip Labs Pvt. Ltd.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from typing import Dict
+
+from dcs_core.core.datasource.base import DataSource
+
+
+class FileDataSource(DataSource):
+    """
+    Abstract class for File data sources
+    """
+
+    def __init__(self, data_source_name: str, data_connection: Dict):
+        super().__init__(data_source_name, data_connection)

{dcs_sdk-1.6.5 → dcs_sdk-1.6.7}/dcs_core/core/datasource/manager.py

@@ -11,6 +11,20 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
+
+#  Copyright 2022-present, the Waterdip Labs Pvt. Ltd.manager
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
 import importlib
 from dataclasses import asdict
 from typing import Dict, List
@@ -43,6 +57,7 @@ class DataSourceManager:
         "oracle": "OracleDataSource",
         "db2": "DB2DataSource",
         "sybase": "SybaseDataSource",
+        "azure_blob": "AzureBlobDataSource",
     }
 
     def __init__(self, config: Configuration):

dcs_sdk-1.6.7/dcs_core/integrations/databases/azure_blob.py

@@ -0,0 +1,115 @@
+#  Copyright 2022-present, the Waterdip Labs Pvt. Ltd.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import io
+from typing import Any, Dict, List, Optional
+
+import pandas as pd
+from azure.storage.blob import BlobServiceClient
+
+from dcs_core.core.common.errors import (
+    DatachecksColumnFetchError,
+    DataChecksDataSourcesConnectionError,
+    DatachecksTableFetchError,
+)
+from dcs_core.core.datasource.file_datasource import FileDataSource
+
+
+class AzureBlobDataSource(FileDataSource):
+    def __init__(self, data_source_name: str, data_connection: Dict):
+        super().__init__(data_source_name, data_connection)
+        self.allowed_file_extensions = [".csv"]
+        self.blob_service_client: Optional[BlobServiceClient] = None
+        self.connection = None
+
+    def connect(self) -> Any:
+        """
+        Connect to the file data source
+        """
+        try:
+            account_name = self.data_connection.get("account_name")
+            container_name = self.data_connection.get("container_name")
+            account_key = self.data_connection.get("account_key")
+            endpoint_suffix = self.data_connection.get("endpoint_suffix", "core.windows.net")
+            connection_str = f"https://{account_name}.blob.{endpoint_suffix}"
+            blob_service_client = BlobServiceClient(account_url=connection_str, credential=account_key)
+            self.blob_service_client = blob_service_client
+            self.connection = blob_service_client.get_container_client(container=container_name)
+            return self.connection
+        except Exception as e:
+            raise DataChecksDataSourcesConnectionError(f"Failed to connect to Azure Blob Storage: {e}")
+
+    def is_connected(self) -> bool:
+        """
+        Check if the file data source is connected
+        """
+        return self.connection is not None
+
+    def close(self):
+        """
+        Close the connection
+        """
+        self.connection.close()
+        self.blob_service_client.close()
+        self.connection = None
+        self.blob_service_client = None
+
+    def query_get_table_names(self) -> dict:
+        """
+        Query to get table names (blob names in this case)
+        """
+        if not self.is_connected():
+            raise DataChecksDataSourcesConnectionError("Not connected to Azure Blob Storage")
+        try:
+            subfolder = self.data_connection.get("subfolder", "")
+            blob_iterator = self.connection.list_blobs(name_starts_with=subfolder)
+            blobs = [
+                blob.name
+                for blob in blob_iterator
+                if len(blob.name.split("/")) == 1 and blob.name.endswith(tuple(self.allowed_file_extensions))
+            ]
+            return {"table": blobs}
+        except Exception as e:
+            raise DatachecksTableFetchError(f"Failed to list blobs: {e}")
+
+    def query_get_table_columns(self, table: str) -> List[dict]:
+        """
+        Get column names for a table (CSV blob in this case).
+        """
+        if not self.is_connected():
+            raise DataChecksDataSourcesConnectionError("Not connected to Azure Blob Storage")
+
+        if not any(table.endswith(ext) for ext in self.allowed_file_extensions):
+            raise ValueError(f"Unsupported file type for {table}. Allowed: {self.allowed_file_extensions}")
+
+        try:
+            blob_client = self.connection.get_blob_client(blob=table)
+            download_stream = blob_client.download_blob()
+            data = download_stream.readall()
+            if table.endswith(".csv"):
+                df = pd.read_csv(io.BytesIO(data))
+            else:
+                raise ValueError(f"Unsupported file type for {table}. Allowed: {self.allowed_file_extensions}")
+
+            return [{"column_name": col, "column_type": "string"} for col in df.columns.tolist()]
+        except Exception as e:
+            raise DatachecksColumnFetchError(f"Failed to read columns from blob '{table}': {e}")
+
+    def query_get_database_version(self) -> str:
+        """
+        Get the database version
+        :return: version string
+        """
+        api_version = self.blob_service_client.api_version
+        return api_version

{dcs_sdk-1.6.5 → dcs_sdk-1.6.7}/dcs_core/integrations/databases/mssql.py

@@ -13,6 +13,7 @@
 #  limitations under the License.
 
 import datetime
+import math
 from decimal import Decimal
 from typing import Any, Dict, List, Optional, Tuple, Union
 from uuid import UUID
@@ -706,13 +707,15 @@ class MssqlDataSource(SQLDataSource):
         cursor = self.connection.cursor()
         try:
             cursor.execute(query)
-            columns = [column[0] for column in cursor.description]
-            result_row = cursor.fetchone()
+            if cursor.description:
+                columns = [column[0] for column in cursor.description]
+                result_row = cursor.fetchone()
+                row = dict(zip(columns, result_row)) if result_row else {}
+            else:
+                row = {}
         finally:
             cursor.close()
 
-        row = dict(zip(columns, result_row))
-
         def _normalize_metrics(value):
             """Safely normalize DB metric values for JSON serialization."""
             if value is None:
@@ -737,11 +740,158 @@ class MssqlDataSource(SQLDataSource):
             col_metrics = {}
 
             for key, value in row.items():
-                if key.startswith(f"{name}_"):
-                    metric_name = key[len(name) + 1 :]
+                clean_key = key.replace("[", "").replace("]", "")
+                if clean_key.startswith(f"{name}_"):
+                    metric_name = clean_key[len(name) + 1 :]
                     col_metrics[metric_name] = _normalize_metrics(value)
 
             column_wise.append({"column_name": name, "metrics": col_metrics})
+
+        for col_data in column_wise:
+            metrics = col_data["metrics"]
+            distinct_count = metrics.get("distinct")
+            col_name = col_data["column_name"]
+
+            dtype = next(c["data_type"].lower() for c in column_info if c["column_name"] == col_name)
+
+            quoted = self.quote_column(col_name)
+
+            is_dtype_numeric = (
+                True
+                if dtype
+                in (
+                    "int",
+                    "integer",
+                    "bigint",
+                    "smallint",
+                    "tinyint",
+                    "decimal",
+                    "numeric",
+                    "float",
+                    "real",
+                    "money",
+                    "smallmoney",
+                )
+                else False
+            )
+
+            if is_dtype_numeric:
+                col_min = metrics.get("min")
+                col_max = metrics.get("max")
+
+                if col_min is not None and col_max is not None and col_min != col_max:
+                    bucket_count = 20
+                    bucket_size = (float(col_max) - float(col_min)) / bucket_count
+
+                    bucket_queries = []
+                    for i in range(bucket_count):
+                        start = float(col_min) + i * bucket_size
+                        end = float(col_min) + (i + 1) * bucket_size
+
+                        bucket_queries.append(
+                            f"SUM(CASE WHEN {quoted} >= {start} AND {quoted} < {end} THEN 1 ELSE 0 END) AS bucket_{i}"
+                        )
+
+                    bucket_sql = f"SELECT {', '.join(bucket_queries)} FROM {qualified_table}"
+
+                    try:
+                        bucket_result = self.connection.execute(text(bucket_sql)).fetchone()
+                        distribution = []
+
+                        for i in range(bucket_count):
+                            start_raw = float(col_min) + i * bucket_size
+                            end_raw = float(col_min) + (i + 1) * bucket_size
+
+                            if dtype in ("int", "integer", "bigint", "smallint", "tinyint"):
+                                start = math.floor(start_raw)
+                                end = math.ceil(end_raw)
+                            else:
+                                start = round(start_raw, 2)
+                                end = round(end_raw, 2)
+
+                            count = bucket_result[i] if bucket_result and bucket_result[i] is not None else 0
+
+                            distribution.append(
+                                {
+                                    "col_val": f"{start} - {end}",
+                                    "count": count,
+                                }
+                            )
+
+                        metrics["distribution_graph"] = distribution
+
+                    except Exception as e:
+                        print(f"Failed to generate numeric distribution for {col_name}: {e}")
+
+                continue
+
+            if isinstance(distinct_count, (int, float)) and distinct_count <= 20:
+                if dtype in ("text", "ntext", "xml"):
+                    group_expr = f"CAST({quoted} AS NVARCHAR(MAX))"
+                else:
+                    group_expr = quoted
+
+                dist_query = (
+                    f"SELECT {group_expr}, COUNT(*) "
+                    f"FROM {qualified_table} GROUP BY {group_expr} ORDER BY COUNT(*) DESC"
+                )
+
+                try:
+                    dist_cursor = self.connection.cursor()
+                    dist_cursor.execute(dist_query)
+                    dist_result = dist_cursor.fetchall()
+                    dist_cursor.close()
+
+                    distribution = []
+
+                    for r in dist_result:
+                        val = _normalize_metrics(r[0])
+                        distribution.append(
+                            {
+                                "col_val": val,
+                                "count": r[1],
+                            }
+                        )
+
+                    metrics["distribution_graph"] = distribution
+
+                except Exception as e:
+                    print(f"Failed to generate distribution graph for column {col_name}: {e}")
+
+        for col_data in column_wise:
+            metrics = col_data["metrics"]
+            distinct_count = metrics.get("distinct")
+            col_name = col_data["column_name"]
+            dtype = next(c["data_type"].lower() for c in column_info if c["column_name"] == col_name)
+
+            quoted = self.quote_column(col_name)
+
+            is_dtype_numeric = (
+                True
+                if dtype
+                in (
+                    "int",
+                    "integer",
+                    "bigint",
+                    "smallint",
+                    "tinyint",
+                    "decimal",
+                    "numeric",
+                    "float",
+                    "real",
+                    "money",
+                    "smallmoney",
+                )
+                else False
+            )
+
+            formatted_metrics_data = {
+                "general_data": {key: value for key, value in metrics.items() if key != "distribution_graph"},
+                "is_dtype_numeric": is_dtype_numeric,
+                "distribution_data": metrics.get("distribution_graph", []),
+            }
+            col_data["metrics"] = formatted_metrics_data
+
         return column_wise
 
     def fetch_sample_values_from_database(

{dcs_sdk-1.6.5 → dcs_sdk-1.6.7}/dcs_core/integrations/databases/postgres.py

@@ -13,6 +13,7 @@
 #  limitations under the License.
 
 import datetime
+import math
 from decimal import Decimal
 from typing import Any, Dict, List, Optional, Tuple
 from uuid import UUID
@@ -411,9 +412,73 @@ class PostgresDataSource(SQLDataSource):
             col_name = col_data["column_name"]
             dtype = next(c["data_type"].lower() for c in column_info if c["column_name"] == col_name)
 
-            if isinstance(distinct_count, (int, float)) and distinct_count < 20:
-                quoted = self.quote_column(col_name)
+            quoted = self.quote_column(col_name)
+
+            is_dtype_numeric = (
+                True
+                if dtype
+                in (
+                    "int",
+                    "integer",
+                    "bigint",
+                    "smallint",
+                    "decimal",
+                    "numeric",
+                    "float",
+                    "double",
+                )
+                else False
+            )
+
+            if is_dtype_numeric:
+                col_min = metrics.get("min")
+                col_max = metrics.get("max")
+
+                if col_min is not None and col_max is not None and col_min != col_max:
+                    bucket_count = 20
+                    bucket_size = (col_max - col_min) / bucket_count
+
+                    bucket_queries = []
+                    for i in range(bucket_count):
+                        start = col_min + i * bucket_size
+                        end = col_min + (i + 1) * bucket_size
+
+                        bucket_queries.append(
+                            f"SUM(CASE WHEN {quoted} >= {start} AND {quoted} < {end} THEN 1 ELSE 0 END) AS bucket_{i}"
+                        )
+
+                    bucket_sql = f"SELECT {', '.join(bucket_queries)} FROM {qualified_table}"
+
+                    try:
+                        bucket_result = self.connection.execute(text(bucket_sql)).fetchone()
+                        distribution = []
+
+                        for i in range(bucket_count):
+                            start_raw = col_min + i * bucket_size
+                            end_raw = col_min + (i + 1) * bucket_size
+                            if dtype in ("int", "integer", "bigint", "smallint"):
+                                start = math.floor(start_raw)
+                                end = math.ceil(end_raw)
+                            else:
+                                start = round(start_raw, 2)
+                                end = round(end_raw, 2)
+                            count = bucket_result[i]
+
+                            distribution.append(
+                                {
+                                    "col_val": f"{start} - {end}",
+                                    "count": count,
+                                }
+                            )
 
+                        metrics["distribution_graph"] = distribution
+
+                    except Exception as e:
+                        print(f"Failed to generate numeric distribution for {col_name}: {e}")
+
+                continue
+
+            if isinstance(distinct_count, (int, float)) and distinct_count <= 20:
                 if dtype in ("json", "jsonb"):
                     group_expr = f"{quoted}::text"
                 else:
@@ -444,8 +509,31 @@ class PostgresDataSource(SQLDataSource):
 
         for col_data in column_wise:
             metrics = col_data["metrics"]
+            distinct_count = metrics.get("distinct")
+            col_name = col_data["column_name"]
+            dtype = next(c["data_type"].lower() for c in column_info if c["column_name"] == col_name)
+
+            quoted = self.quote_column(col_name)
+
+            is_dtype_numeric = (
+                True
+                if dtype
+                in (
+                    "int",
+                    "integer",
+                    "bigint",
+                    "smallint",
+                    "decimal",
+                    "numeric",
+                    "float",
+                    "double",
+                )
+                else False
+            )
+
             formatted_metrics_data = {
                 "general_data": {key: value for key, value in metrics.items() if key != "distribution_graph"},
+                "is_dtype_numeric": is_dtype_numeric,
                 "distribution_data": metrics.get("distribution_graph", []),
             }
             col_data["metrics"] = formatted_metrics_data

{dcs_sdk-1.6.5 → dcs_sdk-1.6.7}/dcs_sdk/__version__.py

@@ -12,4 +12,4 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-__version__ = "1.6.4"
+__version__ = "1.6.7"

{dcs_sdk-1.6.5 → dcs_sdk-1.6.7}/dcs_sdk/sdk/config/config_loader.py

@@ -47,6 +47,7 @@ class SourceTargetConnection(BaseModel):
     port: Optional[Union[int, str]] = None
     driver: str
     table: Optional[str] = None
+    datasource_type: Optional[str] = None
     database: Optional[str] = None
     filepath: Optional[str] = None
     catalog: Optional[str] = None
@@ -66,6 +67,11 @@ class SourceTargetConnection(BaseModel):
     impersonate_service_account: Optional[str] = None  # bigquery specific
     bigquery_credentials: Optional[str] = None  # bigquery specific
     transform_columns: Dict[str, str] | None = None
+    account_name: Optional[str] = None
+    container_name: Optional[str] = None
+    account_key: Optional[str] = None
+    endpoint_suffix: Optional[str] = None
+    subfolder_path: Optional[str] = None
 
 
 class SimilarityConfig(BaseModel):
@@ -140,6 +146,7 @@ class DataDiffConfig:
         "mysql": "mysql",
         "sybase": "sybase",
         "bigquery": "bigquery",
+        "azure_blob": "duckdb",
     }
 
     def __init__(
@@ -307,6 +314,12 @@ class DataDiffConfig:
             "impersonate_service_account": connection.get("connection", {}).get("impersonate_service_account"),
             "bigquery_credentials": connection.get("connection", {}).get("bigquery_credentials"),
             "transform_columns": transform_columns,
+            "datasource_type": connection.get("type"),
+            "account_name": connection.get("connection", {}).get("account_name"),
+            "container_name": connection.get("connection", {}).get("container_name"),
+            "account_key": connection.get("connection", {}).get("account_key"),
+            "endpoint_suffix": connection.get("connection", {}).get("endpoint_suffix"),
+            "subfolder_path": connection.get("connection", {}).get("subfolder_path"),
         }
 
     def get_data_diff_configs(self) -> List[Comparison]: