dcs-sdk 1.7.4__py3-none-any.whl → 1.7.5__py3-none-any.whl

dcs_core/integrations/databases/databricks.py

@@ -12,12 +12,17 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from typing import Any, Dict
+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
+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
 
 
@@ -44,8 +49,505 @@ class DatabricksDataSource(SQLDataSource):
             )
             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/__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.dist-info → dcs_sdk-1.7.5.dist-info}/METADATA

@@ -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.4.dist-info → dcs_sdk-1.7.5.dist-info}/RECORD

@@ -102,7 +102,7 @@ dcs_core/integrations/__init__.py,sha256=RkfhRKLXEForLCs4rZkTf0qc_b0TokSggSAcKI4
 dcs_core/integrations/databases/__init__.py,sha256=RkfhRKLXEForLCs4rZkTf0qc_b0TokSggSAcKI4yfZg,610
 dcs_core/integrations/databases/azure_blob.py,sha256=XF-B790XA1sGyhgIWUPE-BRHR_-ctA31IjI66pNL6eM,4231
 dcs_core/integrations/databases/bigquery.py,sha256=26RuypLMmiARZIWkV_mxtnNL2yCs94YWerSGH5Nr10Q,7337
-dcs_core/integrations/databases/databricks.py,sha256=n4fm5m_mtRCdtjLGDvbNW18u7Ev234vDBjq_lxuOxns,1978
+dcs_core/integrations/databases/databricks.py,sha256=r27uY_XtlJdbgeQ57ABbcfV6tU66LnbWyQucStayAVo,20768
 dcs_core/integrations/databases/db2.py,sha256=hNGivvYCitp88ouZlCxp7iRQ-vnPiK1kL8x85NyGotk,26492
 dcs_core/integrations/databases/duck_db.py,sha256=X4FRSsobOFCIi329cYofQsMd_fkRI4KxC8BIrtiDz4g,5531
 dcs_core/integrations/databases/elasticsearch.py,sha256=6CTGs1WGrfgdDRNVt9DpOB0_z_znT6YoVj10E1WY-wQ,2152
@@ -134,7 +134,7 @@ dcs_core/report/static/index.js,sha256=p4wvku-zlXi0y4gWeSzV1amY0s4mjtUq2QsezARLV
 dcs_core/report/static/index.js.LICENSE.txt,sha256=bBDZBJVEDrqjCi7sfoF8CchjFn3hdcbNkP7ub7kbcXQ,201041
 dcs_sdk/__init__.py,sha256=RkfhRKLXEForLCs4rZkTf0qc_b0TokSggSAcKI4yfZg,610
 dcs_sdk/__main__.py,sha256=Qn8stIaQGrdLjHQ-H7xO0T-brtq5RWZoWU9QvqoarV8,683
-dcs_sdk/__version__.py,sha256=ZRK4QBUilILjo7h5_d7FuTLhbBM4aJj6kS5UDIrdZH4,633
+dcs_sdk/__version__.py,sha256=LyQoYvx5NK-5Cr6vd2YXpSX8fYqK7p67ywpDOwSpWEI,633
 dcs_sdk/cli/__init__.py,sha256=RkfhRKLXEForLCs4rZkTf0qc_b0TokSggSAcKI4yfZg,610
 dcs_sdk/cli/cli.py,sha256=jaO52UrMWLafcF_yhqllPkmYSTuO2sksFi30fYFdAB4,4406
 dcs_sdk/sdk/__init__.py,sha256=skrZcgWWJBL6NXTUERywJ3qRJRemgpDXyW7lPg1FJk8,2107
@@ -156,7 +156,7 @@ dcs_sdk/sdk/utils/similarity_score/levenshtein_distance_provider.py,sha256=puAWP
 dcs_sdk/sdk/utils/table.py,sha256=X8HxdYTWyx_oVrBWPsXlmA-xJKXXDBW9RrhlWNqA1As,18224
 dcs_sdk/sdk/utils/themes.py,sha256=Meo2Yldv4uyPpEqI7qdA28Aa6sxtwUU1dLKKm4QavjM,1403
 dcs_sdk/sdk/utils/utils.py,sha256=F7BUnLgJFyb7mc_9cw1hyk4lutsg6ctJuBQRrfyBv0g,16874
-dcs_sdk-1.7.4.dist-info/METADATA,sha256=qVNFesr-vlH-JvTyy9haJemG2TK-lQURlJiNtX1q2Ls,7652
-dcs_sdk-1.7.4.dist-info/WHEEL,sha256=3ny-bZhpXrU6vSQ1UPG34FoxZBp3lVcvK0LkgUz6VLk,88
-dcs_sdk-1.7.4.dist-info/entry_points.txt,sha256=XhODNz7UccgPOyklXgp7pIfTTXArd6-V0mImjhnhwto,80
-dcs_sdk-1.7.4.dist-info/RECORD,,
+dcs_sdk-1.7.5.dist-info/METADATA,sha256=VLsUM65db7o2L5J6UpwF5yU2_bz3Pmx13sYQOZknY7U,7652
+dcs_sdk-1.7.5.dist-info/WHEEL,sha256=kJCRJT_g0adfAJzTx2GUMmS80rTJIVHRCfG0DQgLq3o,88
+dcs_sdk-1.7.5.dist-info/entry_points.txt,sha256=XhODNz7UccgPOyklXgp7pIfTTXArd6-V0mImjhnhwto,80
+dcs_sdk-1.7.5.dist-info/RECORD,,

{dcs_sdk-1.7.4.dist-info → dcs_sdk-1.7.5.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 2.3.0
+Generator: poetry-core 2.3.1
 Root-Is-Purelib: true
 Tag: py3-none-any