dcs-sdk 1.7.4__tar.gz → 1.7.5__tar.gz

{dcs_sdk-1.7.4 → dcs_sdk-1.7.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dcs-sdk
-Version: 1.7.4
+Version: 1.7.5
 Summary: SDK for DataChecks
 Author: Waterdip Labs
 Author-email: hello@waterdip.ai

dcs_sdk-1.7.5/dcs_core/integrations/databases/databricks.py

@@ -0,0 +1,553 @@
+#  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 datetime
+import math
+from decimal import Decimal
+from typing import Any, Dict, List, Optional, Tuple
+from uuid import UUID
+
+from sqlalchemy import create_engine, text
+from sqlalchemy.engine import URL
+
+from dcs_core.core.common.errors import DataChecksDataSourcesConnectionError
+from dcs_core.core.common.models.data_source_resource import RawColumnInfo
+from dcs_core.core.datasource.sql_datasource import SQLDataSource
+
+
+class DatabricksDataSource(SQLDataSource):
+    def __init__(self, data_source_name: str, data_connection: Dict):
+        super().__init__(data_source_name, data_connection)
+
+    def connect(self) -> Any:
+        """
+        Connect to the data source
+        """
+        try:
+            url = URL.create(
+                "databricks",
+                username="token",
+                password=self.data_connection.get("token"),
+                host=self.data_connection.get("host"),
+                port=self.data_connection.get("port", 443),
+                database=self.data_connection.get("schema"),
+                query={
+                    "http_path": self.data_connection.get("http_path"),
+                    "catalog": self.data_connection.get("catalog"),
+                },
+            )
+            engine = create_engine(url, echo=True)
+            self.connection = engine.connect()
+            self.schema_name = self.data_connection.get("schema")
+            return self.connection
+        except Exception as e:
+            raise DataChecksDataSourcesConnectionError(
+                message=f"Failed to connect to Databricks data source: [{str(e)}]"
+            )
+
+    def quote_column(self, column: str) -> str:
+        return f"`{column}`"
+
+    def quote_database(self, database: str) -> str:
+        return f"`{database}`"
+
+    def qualified_table_name(self, table_name: str) -> str:
+        if self.schema_name:
+            return f"`{self.schema_name}`.`{table_name}`"
+        return f"`{table_name}`"
+
+    def query_get_database_version(self, database_version_query: Optional[str] = None) -> str:
+        """
+        Get the database version
+        :return: version string
+        """
+        query = database_version_query or "SELECT version()"
+        result = self.fetchone(query)
+        if result:
+            return result[0]
+        return None
+
+    def query_get_table_indexes(self, table: str, schema: str | None = None) -> dict[str, dict]:
+        """
+        Get index information for a table.
+        For Databricks, this primarily returns Primary Key information as traditional indexes are not exposed identically.
+        :param table: Table name
+        :param schema: Optional schema name
+        :return: Dictionary with index details
+        """
+        schema = schema or self.schema_name
+        database = self.data_connection.get("catalog") or "hive_metastore"
+        quoted_database = self.quote_database(database)
+
+        # Databricks Unity Catalog stores constraints in information_schema
+        # We will fetch Primary Key info and structure it as an "index"
+        query = f"""
+            SELECT
+                tc.constraint_name,
+                kcu.column_name,
+                kcu.ordinal_position
+            FROM {quoted_database}.information_schema.table_constraints AS tc
+            JOIN {quoted_database}.information_schema.key_column_usage AS kcu
+                ON tc.constraint_name = kcu.constraint_name
+                AND tc.table_schema = kcu.table_schema
+            WHERE tc.table_schema = '{schema}'
+              AND tc.table_name = '{table}'
+              AND tc.constraint_type = 'PRIMARY KEY'
+            ORDER BY kcu.ordinal_position
+        """
+
+        try:
+            rows = self.fetchall(query)
+        except Exception as e:
+            # Fallback or silent failure if table doesn't exist or info schema not accessible
+            return {}
+
+        indexes = {}
+        if rows:
+            # In Databricks, the PK constraint name acts as the index name for this purpose
+            constraint_name = rows[0][0]
+            indexes[constraint_name] = {"columns": [], "index_type": "PRIMARY KEY", "is_primary_key": True}
+
+            for row in rows:
+                col_name = row[1]
+                ordinal = row[2]
+                indexes[constraint_name]["columns"].append({"column_name": col_name, "column_order": ordinal})
+
+        return indexes
+
+    def get_table_foreign_key_info(self, table_name: str, schema: str | None = None) -> list[dict]:
+        """
+        Get foreign key information for a table.
+        :param table_name: Table name
+        :param schema: Optional schema name
+        :return: List of dicts with FK details
+        """
+        schema = schema or self.schema_name
+        database = self.data_connection.get("catalog") or "hive_metastore"
+        quoted_database = self.quote_database(database)
+
+        # Standard ISO SQL query for Foreign Keys using information_schema
+        # Works for Unity Catalog
+        query = f"""
+            SELECT
+                tc.constraint_name,
+                tc.table_name,
+                kcu.column_name AS fk_column,
+                rel_kcu.table_name AS referenced_table,
+                rel_kcu.column_name AS referenced_column
+            FROM {quoted_database}.information_schema.table_constraints tc
+            JOIN {quoted_database}.information_schema.key_column_usage kcu
+                ON tc.constraint_name = kcu.constraint_name
+                AND tc.table_schema = kcu.table_schema
+            JOIN {quoted_database}.information_schema.referential_constraints rc
+                ON tc.constraint_name = rc.constraint_name
+                AND tc.table_schema = rc.constraint_schema
+            JOIN {quoted_database}.information_schema.key_column_usage rel_kcu
+                ON rc.unique_constraint_name = rel_kcu.constraint_name
+                AND rc.unique_constraint_schema = rel_kcu.table_schema
+                AND kcu.ordinal_position = rel_kcu.ordinal_position
+            WHERE tc.constraint_type = 'FOREIGN KEY'
+              AND tc.table_name = '{table_name}'
+              AND tc.table_schema = '{schema}'
+        """
+
+        try:
+            rows = self.fetchall(query)
+        except Exception:
+            return []
+
+        fk_info = []
+        for row in rows:
+            fk_info.append(
+                {
+                    "constraint_name": row[0],
+                    "table_name": row[1],
+                    "fk_column": row[2],
+                    "referenced_table": row[3],
+                    "referenced_column": row[4],
+                }
+            )
+
+        return fk_info
+
+    def query_get_table_names(
+        self,
+        schema: str | None = None,
+        with_view: bool = False,
+    ) -> dict:
+        """
+        Get the list of tables in the database.
+        :param schema: optional schema name
+        :param with_view: whether to include views
+        :return: dictionary with table names and optionally view names
+        """
+        schema = schema or self.schema_name
+        database = self.data_connection.get("catalog") or "hive_metastore"
+        quoted_database = self.quote_database(database)
+
+        if with_view:
+            table_type_condition = "table_type IN ('MANAGED', 'EXTERNAL', 'VIEW')"
+        else:
+            table_type_condition = "table_type IN ('MANAGED', 'EXTERNAL')"
+
+        query = (
+            f"SELECT table_name, table_type FROM {quoted_database}.information_schema.tables "
+            f"WHERE table_schema = '{schema}' "
+        )
+        if not with_view:
+            query += " AND table_type != 'VIEW'"
+
+        rows = self.fetchall(query)
+
+        if with_view:
+            result = {"table": [], "view": []}
+            if rows:
+                for row in rows:
+                    table_name = row[0]
+                    table_type = row[1]
+                    if "VIEW" in table_type:
+                        result["view"].append(table_name)
+                    else:
+                        result["table"].append(table_name)
+        else:
+            result = {"table": []}
+            if rows:
+                result["table"] = [row[0] for row in rows]
+
+        return result
+
+    def query_get_table_columns(
+        self,
+        table: str,
+        schema: str | None = None,
+    ) -> RawColumnInfo:
+        """
+        Get the schema of a table.
+        :param table: table name
+        :return: RawColumnInfo object containing column information
+        """
+        schema = schema or self.schema_name
+        database = self.data_connection.get("catalog") or "hive_metastore"
+        quoted_database = self.quote_database(database)
+
+        query = (
+            f"SELECT column_name, data_type, datetime_precision, numeric_precision, numeric_scale, "
+            f"character_maximum_length "
+            f"FROM {quoted_database}.information_schema.columns "
+            f"WHERE table_name = '{table}' AND table_schema = '{schema}'"
+        )
+        rows = self.fetchall(query)
+
+        if not rows:
+            raise RuntimeError(f"{table}: Table, {schema}: Schema, does not exist, or has no columns")
+
+        column_info = {
+            r[0]: RawColumnInfo(
+                column_name=self.safe_get(r, 0),
+                data_type=self.safe_get(r, 1),
+                datetime_precision=self.safe_get(r, 2),
+                numeric_precision=self.safe_get(r, 3),
+                numeric_scale=self.safe_get(r, 4),
+                character_maximum_length=self.safe_get(r, 5),
+            )
+            for r in rows
+        }
+        return column_info
+
+    def fetch_rows(
+        self,
+        query: str,
+        limit: int = 1,
+        with_column_names: bool = False,
+        complete_query: Optional[str] = None,
+    ) -> Tuple[List, Optional[List[str]]]:
+        """
+        Fetch rows from the database.
+
+        :param query: SQL query to execute.
+        :param limit: Number of rows to fetch.
+        :param with_column_names: Whether to include column names in the result.
+        :return: Tuple of (rows, column_names or None)
+        """
+        query = complete_query or f"SELECT * FROM ({query}) AS subquery LIMIT {limit}"
+
+        result = self.connection.execute(text(query))
+        rows = result.fetchmany(limit)
+
+        if with_column_names:
+            column_names = result.keys()
+            return rows, list(column_names)
+        else:
+            return rows, None
+
+    def fetch_sample_values_from_database(
+        self,
+        table_name: str,
+        column_names: list[str],
+        limit: int = 5,
+    ) -> List[Tuple]:
+        """
+        Fetch sample rows for specific columns from the given table.
+
+        :param table_name: The name of the table.
+        :param column_names: List of column names to fetch.
+        :param limit: Number of rows to fetch.
+        :return: List of row tuples.
+        """
+        table_name = self.qualified_table_name(table_name)
+
+        if not column_names:
+            raise ValueError("At least one column name must be provided")
+
+        if len(column_names) == 1 and column_names[0] == "*":
+            query = f"SELECT * FROM {table_name} LIMIT {limit}"
+        else:
+            columns = ", ".join([self.quote_column(col) for col in column_names])
+            query = f"SELECT {columns} FROM {table_name} LIMIT {limit}"
+
+        result = self.connection.execute(text(query))
+        column_names = list(result.keys())
+        rows = result.fetchall()
+        return rows, column_names
+
+    def build_table_metrics_query(
+        self,
+        table_name: str,
+        column_info: list[dict],
+        additional_queries: Optional[List[str]] = None,
+    ) -> list[dict]:
+        query_parts = []
+        if not column_info:
+            return []
+
+        for col in column_info:
+            name = col["column_name"]
+            dtype = col["data_type"].lower()
+            quoted = self.quote_column(name)
+
+            if dtype in ("string", "varchar"):
+                distinct_expr = f"{quoted}"
+            else:
+                distinct_expr = f"{quoted}"
+
+            query_parts.append(f"COUNT(DISTINCT {distinct_expr}) AS `{name}_distinct`")
+            query_parts.append(f"COUNT(*) - COUNT(DISTINCT {distinct_expr}) AS `{name}_duplicate`")
+            query_parts.append(
+                f"SUM(CASE WHEN {self.quote_column(name)} IS NULL THEN 1 ELSE 0 END) AS `{name}_is_null`"
+            )
+
+            if dtype in (
+                "int",
+                "integer",
+                "bigint",
+                "long",
+                "smallint",
+                "tinyint",
+                "decimal",
+                "numeric",
+                "float",
+                "double",
+            ):
+                query_parts.append(f"MIN({self.quote_column(name)}) AS `{name}_min`")
+                query_parts.append(f"MAX({self.quote_column(name)}) AS `{name}_max`")
+                query_parts.append(f"AVG({self.quote_column(name)}) AS `{name}_average`")
+
+            elif dtype in ("string", "varchar", "char", "text"):
+                # Databricks uses length() or char_length()
+                query_parts.append(f"MAX(length({self.quote_column(name)})) AS `{name}_max_character_length`")
+
+            elif dtype in ("timestamp", "date", "boolean"):
+                query_parts.append(f"MIN({self.quote_column(name)}) AS `{name}_min`")
+                query_parts.append(f"MAX({self.quote_column(name)}) AS `{name}_max`")
+
+        if additional_queries:
+            for queries in additional_queries:
+                query_parts.append(queries)
+
+        qualified_table = self.qualified_table_name(table_name)
+        joined_parts = ",\n    ".join(query_parts)
+        query = f"SELECT\n    {joined_parts}\nFROM {qualified_table};"
+
+        result = self.connection.execute(text(query))
+        row = dict(list(result)[0]._mapping)
+
+        def _normalize_metrics(value):
+            """
+            Safely normalizes DB metric values into JSON-serializable Python types.
+            """
+            if value is None:
+                return None
+
+            if isinstance(value, Decimal):
+                return float(value)
+            if isinstance(value, (int, float, bool)):
+                return value
+
+            if isinstance(value, (datetime.datetime, datetime.date)):
+                return value.isoformat()
+
+            if isinstance(value, UUID):
+                return str(value)
+
+            if isinstance(value, list):
+                return [_normalize_metrics(v) for v in value]
+            if isinstance(value, dict):
+                return {k: _normalize_metrics(v) for k, v in value.items()}
+
+            return str(value)
+
+        column_wise = []
+        for col in column_info:
+            name = col["column_name"]
+            col_metrics = {}
+
+            for key, value in row.items():
+                if key.startswith(f"{name}_"):
+                    metric_name = 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",
+                    "long",
+                    "smallint",
+                    "tinyint",
+                    "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
+
+                        # Databricks SQL syntax for CASE WHEN
+                        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", "long", "smallint", "tinyint"):
+                                start = math.floor(start_raw)
+                                end = math.ceil(end_raw)
+                            else:
+                                start = round(start_raw, 2)
+                                end = round(end_raw, 2)
+
+                            # Fetch by index or name (sqlalchemy row access)
+                            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:
+                group_expr = quoted
+
+                dist_query = (
+                    f"SELECT {group_expr}, COUNT(*) "
+                    f"FROM {qualified_table} GROUP BY {group_expr} ORDER BY COUNT(*) DESC"
+                )
+
+                try:
+                    dist_result = self.connection.execute(text(dist_query)).fetchall()
+
+                    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"]
+            # Formatting as per existing pattern
+            is_dtype_numeric = (
+                True
+                if next(c["data_type"].lower() for c in column_info if c["column_name"] == col_data["column_name"])
+                in (
+                    "int",
+                    "integer",
+                    "bigint",
+                    "long",
+                    "smallint",
+                    "tinyint",
+                    "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
+
+        return column_wise

{dcs_sdk-1.7.4 → dcs_sdk-1.7.5}/dcs_sdk/__version__.py

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

{dcs_sdk-1.7.4 → dcs_sdk-1.7.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "dcs-sdk"
-version = "1.7.4"
+version = "1.7.5"
 description = "SDK for DataChecks"
 authors = ["Waterdip Labs <hello@waterdip.ai>"]
 readme = "README.md"

dcs_sdk-1.7.4/dcs_core/integrations/databases/databricks.py

@@ -1,51 +0,0 @@
-#  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 Any, Dict
-
-from sqlalchemy import create_engine
-from sqlalchemy.engine import URL
-
-from dcs_core.core.common.errors import DataChecksDataSourcesConnectionError
-from dcs_core.core.datasource.sql_datasource import SQLDataSource
-
-
-class DatabricksDataSource(SQLDataSource):
-    def __init__(self, data_source_name: str, data_connection: Dict):
-        super().__init__(data_source_name, data_connection)
-
-    def connect(self) -> Any:
-        """
-        Connect to the data source
-        """
-        try:
-            url = URL.create(
-                "databricks",
-                username="token",
-                password=self.data_connection.get("token"),
-                host=self.data_connection.get("host"),
-                port=self.data_connection.get("port", 443),
-                database=self.data_connection.get("schema"),
-                query={
-                    "http_path": self.data_connection.get("http_path"),
-                    "catalog": self.data_connection.get("catalog"),
-                },
-            )
-            engine = create_engine(url, echo=True)
-            self.connection = engine.connect()
-            return self.connection
-        except Exception as e:
-            raise DataChecksDataSourcesConnectionError(
-                message=f"Failed to connect to Databricks data source: [{str(e)}]"
-            )