aptdata 0.0.2__py3-none-any.whl

aptdata/plugins/postgres.py

@@ -0,0 +1,113 @@
+"""PostgreSQL reader / writer plugin.
+
+Provides :class:`PostgresReader` and :class:`PostgresWriter` for
+interacting with PostgreSQL databases via **SQLAlchemy**.
+
+Both ``sqlalchemy`` and a PostgreSQL driver (e.g. ``psycopg2-binary``)
+are required.  A friendly
+:class:`~aptdata.plugins.manager.PluginDependencyError` is raised
+when either is missing.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from aptdata.core.dataset import BaseDataset
+from aptdata.plugins.base import BaseReader, BaseWriter
+from aptdata.plugins.dataset import InMemoryDataset
+from aptdata.plugins.manager import PluginDependencyError
+
+
+def _require_sqlalchemy() -> Any:
+    """Import and return the ``sqlalchemy`` module, or raise a friendly error."""
+    try:
+        import sqlalchemy  # noqa: WPS433
+    except ImportError:
+        raise PluginDependencyError("postgres", "sqlalchemy") from None
+    return sqlalchemy
+
+
+class PostgresReader(BaseReader):
+    """Execute a SQL query against a PostgreSQL database and return the result
+    as an :class:`~aptdata.plugins.dataset.InMemoryDataset`.
+
+    Parameters
+    ----------
+    connection_url:
+        SQLAlchemy connection URL, e.g.
+        ``"postgresql+psycopg2://user:pass@host:5432/db"``.
+    query:
+        Raw SQL ``SELECT`` query to execute.
+    """
+
+    def __init__(self, connection_url: str, query: str) -> None:
+        self.connection_url = connection_url
+        self.query = query
+
+    def read(self, **kwargs: Any) -> InMemoryDataset:
+        sa = _require_sqlalchemy()
+        engine = sa.create_engine(self.connection_url)
+        with engine.connect() as conn:
+            result = conn.execute(sa.text(self.query))
+            columns = list(result.keys())
+            records = [dict(zip(columns, row)) for row in result.fetchall()]
+
+        ds = InMemoryDataset(uri=self.connection_url)
+        ds.write(records)
+        return ds
+
+
+class PostgresWriter(BaseWriter):
+    """Write an :class:`~aptdata.plugins.dataset.InMemoryDataset` to a
+    PostgreSQL table.
+
+    Parameters
+    ----------
+    connection_url:
+        SQLAlchemy connection URL.
+    table:
+        Target table name.
+    if_exists:
+        Behaviour when the table already exists: ``"append"`` (default),
+        ``"replace"``, or ``"fail"``.
+    """
+
+    def __init__(
+        self,
+        connection_url: str,
+        table: str,
+        *,
+        if_exists: str = "append",
+    ) -> None:
+        self.connection_url = connection_url
+        self.table = table
+        self.if_exists = if_exists
+
+    def write(self, dataset: BaseDataset, **kwargs: Any) -> None:
+        sa = _require_sqlalchemy()
+        records: list[dict[str, Any]] = dataset.read()
+        if not records:
+            return
+
+        engine = sa.create_engine(self.connection_url)
+        meta = sa.MetaData()
+
+        with engine.connect() as conn:
+            if self.if_exists == "replace":
+                # Use SQLAlchemy DDL to avoid raw SQL interpolation
+                tbl = sa.Table(self.table, sa.MetaData())
+                tbl.drop(engine, checkfirst=True)
+
+            # Auto-create a simple text-column table when it doesn't exist
+            if not sa.inspect(engine).has_table(self.table):
+                columns = [sa.Column(k, sa.Text) for k in records[0]]
+                sa.Table(self.table, meta, *columns)
+                meta.create_all(engine)
+
+            table_obj = sa.Table(self.table, sa.MetaData(), autoload_with=engine)
+            conn.execute(table_obj.insert(), records)
+            conn.commit()
+
+
+__all__ = ["PostgresReader", "PostgresWriter"]

aptdata/plugins/quality/__init__.py

@@ -0,0 +1,39 @@
+"""Data quality plugin package.
+
+Re-exports the main public API for data contracts, quality expectations,
+validation, and reporting.
+"""
+
+from __future__ import annotations
+
+from aptdata.plugins.quality.contract import (
+    ColumnClassification,
+    ColumnContract,
+    EnforcementMode,
+    SchemaContract,
+)
+from aptdata.plugins.quality.expectations import (
+    BaseExpectation,
+    ExpectColumnToNotBeNull,
+    ExpectColumnValuesInRange,
+    ExpectColumnValuesToBeUnique,
+    ExpectColumnValuesToMatchRegex,
+)
+from aptdata.plugins.quality.report import CheckResult, CheckStatus, QualityReport
+from aptdata.plugins.quality.validator import QualityValidator
+
+__all__ = [
+    "ColumnClassification",
+    "ColumnContract",
+    "EnforcementMode",
+    "SchemaContract",
+    "BaseExpectation",
+    "ExpectColumnToNotBeNull",
+    "ExpectColumnValuesInRange",
+    "ExpectColumnValuesToBeUnique",
+    "ExpectColumnValuesToMatchRegex",
+    "CheckResult",
+    "CheckStatus",
+    "QualityReport",
+    "QualityValidator",
+]

aptdata/plugins/quality/contract.py

@@ -0,0 +1,128 @@
+"""Schema contracts and column classification definitions.
+
+Provides :class:`SchemaContract` for declaring the expected schema of a
+dataset, including column types, nullability, classification, and PII
+annotations.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+
+class ColumnClassification(str, Enum):
+    """Data sensitivity classification for a column or dataset."""
+
+    PUBLIC = "PUBLIC"
+    INTERNAL = "INTERNAL"
+    CONFIDENTIAL = "CONFIDENTIAL"
+    PII = "PII"
+    PHI = "PHI"
+    FINANCIAL = "FINANCIAL"
+    SENSITIVE = "SENSITIVE"
+
+
+class EnforcementMode(str, Enum):
+    """How the quality framework reacts when a contract is violated.
+
+    ABORT
+        Raise an exception immediately.
+    WARN
+        Log a warning but continue processing.
+    TAG
+        Annotate the data with quality metadata and continue.
+    """
+
+    ABORT = "ABORT"
+    WARN = "WARN"
+    TAG = "TAG"
+
+
+@dataclass
+class ColumnContract:
+    """Contract specification for a single column.
+
+    Parameters
+    ----------
+    name:
+        Column name as it appears in the dataset.
+    dtype:
+        Expected data type (e.g. ``"int64"``, ``"str"``).
+    nullable:
+        Whether ``null`` / ``None`` values are allowed.
+    classification:
+        Sensitivity classification (see :class:`ColumnClassification`).
+    description:
+        Human-readable description.
+    pii:
+        Shorthand flag indicating the column contains personally
+        identifiable information.
+    retention_days:
+        Number of days the column's data must be retained.
+    metadata:
+        Arbitrary extra metadata.
+    """
+
+    name: str
+    dtype: str = ""
+    nullable: bool = True
+    classification: ColumnClassification = ColumnClassification.INTERNAL
+    description: str = ""
+    pii: bool = False
+    retention_days: int = 0
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class SchemaContract:
+    """Contract specification for an entire dataset schema.
+
+    Parameters
+    ----------
+    name:
+        Contract identifier.
+    version:
+        Semantic version string (e.g. ``"1.0.0"``).
+    owner:
+        Team or person responsible for this contract.
+    description:
+        Human-readable description.
+    columns:
+        Ordered list of :class:`ColumnContract` definitions.
+    enforcement:
+        How violations are handled (see :class:`EnforcementMode`).
+    metadata:
+        Arbitrary extra metadata.
+    """
+
+    name: str
+    version: str = "1.0.0"
+    owner: str = ""
+    description: str = ""
+    columns: list[ColumnContract] = field(default_factory=list)
+    enforcement: EnforcementMode = EnforcementMode.ABORT
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def get_pii_columns(self) -> list[ColumnContract]:
+        """Return columns that are flagged as PII."""
+        return [
+            c
+            for c in self.columns
+            if c.pii or c.classification == ColumnClassification.PII
+        ]
+
+    def get_columns_by_classification(
+        self, classification: ColumnClassification
+    ) -> list[ColumnContract]:
+        """Return columns whose classification matches *classification*."""
+        return [c for c in self.columns if c.classification == classification]
+
+
+__all__ = [
+    "ColumnClassification",
+    "EnforcementMode",
+    "ColumnContract",
+    "SchemaContract",
+]

aptdata/plugins/quality/expectations.py

@@ -0,0 +1,310 @@
+"""Data quality expectations — engine-agnostic validators.
+
+Each :class:`BaseExpectation` automatically dispatches to either
+:meth:`validate_pandas` or :meth:`validate_spark` based on the type of the
+DataFrame passed to :meth:`validate`.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any
+
+from aptdata.plugins.quality.report import CheckResult, CheckStatus
+
+
+def _is_spark_df(df: Any) -> bool:
+    """Return ``True`` when *df* looks like a PySpark DataFrame."""
+    return "pyspark" in type(df).__module__
+
+
+class BaseExpectation(ABC):
+    """Abstract base for all data quality expectations.
+
+    Subclasses implement :meth:`validate_pandas` and :meth:`validate_spark`.
+    The concrete :meth:`validate` method dispatches automatically.
+    """
+
+    @abstractmethod
+    def validate_pandas(self, df: Any) -> CheckResult:
+        """Validate a ``pd.DataFrame`` and return a :class:`CheckResult`."""
+
+    @abstractmethod
+    def validate_spark(self, df: Any) -> CheckResult:
+        """Validate a PySpark ``DataFrame`` and return a :class:`CheckResult`."""
+
+    def validate(self, df: Any) -> CheckResult:
+        """Validate *df* dispatching to the appropriate engine implementation."""
+        if _is_spark_df(df):
+            return self.validate_spark(df)
+        return self.validate_pandas(df)
+
+
+@dataclass
+class ExpectColumnToNotBeNull(BaseExpectation):
+    """Expect that a column contains no null values.
+
+    Parameters
+    ----------
+    column:
+        Name of the column to check.
+    """
+
+    column: str
+
+    def validate_pandas(self, df: Any) -> CheckResult:
+        """Check for null values using pandas."""
+        if self.column not in df.columns:
+            return CheckResult(
+                expectation_name="ExpectColumnToNotBeNull",
+                column=self.column,
+                status=CheckStatus.FAILED,
+                message=f"Column '{self.column}' not found in DataFrame.",
+                rows_evaluated=len(df),
+                rows_failed=len(df),
+            )
+        null_count = int(df[self.column].isnull().sum())
+        status = CheckStatus.PASSED if null_count == 0 else CheckStatus.FAILED
+        return CheckResult(
+            expectation_name="ExpectColumnToNotBeNull",
+            column=self.column,
+            status=status,
+            message=(
+                f"Column '{self.column}' has {null_count} null value(s)."
+                if null_count > 0
+                else f"Column '{self.column}' has no null values."
+            ),
+            rows_evaluated=len(df),
+            rows_failed=null_count,
+        )
+
+    def validate_spark(self, df: Any) -> CheckResult:
+        """Check for null values using PySpark."""
+        from pyspark.sql import functions as F  # noqa: N812
+
+        total = df.count()
+        null_count = df.filter(F.col(self.column).isNull()).count()
+        status = CheckStatus.PASSED if null_count == 0 else CheckStatus.FAILED
+        return CheckResult(
+            expectation_name="ExpectColumnToNotBeNull",
+            column=self.column,
+            status=status,
+            message=(
+                f"Column '{self.column}' has {null_count} null value(s)."
+                if null_count > 0
+                else f"Column '{self.column}' has no null values."
+            ),
+            rows_evaluated=total,
+            rows_failed=null_count,
+        )
+
+
+@dataclass
+class ExpectColumnValuesInRange(BaseExpectation):
+    """Expect that all values in a column fall within [min_val, max_val].
+
+    Parameters
+    ----------
+    column:
+        Name of the column to check.
+    min_val:
+        Inclusive lower bound.
+    max_val:
+        Inclusive upper bound.
+    """
+
+    column: str
+    min_val: float
+    max_val: float
+
+    def validate_pandas(self, df: Any) -> CheckResult:
+        """Check numeric range using pandas."""
+        if self.column not in df.columns:
+            return CheckResult(
+                expectation_name="ExpectColumnValuesInRange",
+                column=self.column,
+                status=CheckStatus.FAILED,
+                message=f"Column '{self.column}' not found in DataFrame.",
+                rows_evaluated=len(df),
+                rows_failed=len(df),
+            )
+        series = df[self.column]
+        out_of_range = int(((series < self.min_val) | (series > self.max_val)).sum())
+        status = CheckStatus.PASSED if out_of_range == 0 else CheckStatus.FAILED
+        return CheckResult(
+            expectation_name="ExpectColumnValuesInRange",
+            column=self.column,
+            status=status,
+            message=(
+                f"Column '{self.column}' has {out_of_range} value(s) outside "
+                f"[{self.min_val}, {self.max_val}]."
+                if out_of_range > 0
+                else f"All values in '{self.column}' are within range."
+            ),
+            rows_evaluated=len(df),
+            rows_failed=out_of_range,
+            metadata={"min_val": self.min_val, "max_val": self.max_val},
+        )
+
+    def validate_spark(self, df: Any) -> CheckResult:
+        """Check numeric range using PySpark."""
+        from pyspark.sql import functions as F  # noqa: N812
+
+        total = df.count()
+        out_of_range = df.filter(
+            (F.col(self.column) < self.min_val) | (F.col(self.column) > self.max_val)
+        ).count()
+        status = CheckStatus.PASSED if out_of_range == 0 else CheckStatus.FAILED
+        return CheckResult(
+            expectation_name="ExpectColumnValuesInRange",
+            column=self.column,
+            status=status,
+            message=(
+                f"Column '{self.column}' has {out_of_range} value(s) outside "
+                f"[{self.min_val}, {self.max_val}]."
+                if out_of_range > 0
+                else f"All values in '{self.column}' are within range."
+            ),
+            rows_evaluated=total,
+            rows_failed=out_of_range,
+            metadata={"min_val": self.min_val, "max_val": self.max_val},
+        )
+
+
+@dataclass
+class ExpectColumnValuesToBeUnique(BaseExpectation):
+    """Expect that all values in a column are unique.
+
+    Parameters
+    ----------
+    column:
+        Name of the column to check.
+    """
+
+    column: str
+
+    def validate_pandas(self, df: Any) -> CheckResult:
+        """Check uniqueness using pandas."""
+        if self.column not in df.columns:
+            return CheckResult(
+                expectation_name="ExpectColumnValuesToBeUnique",
+                column=self.column,
+                status=CheckStatus.FAILED,
+                message=f"Column '{self.column}' not found in DataFrame.",
+                rows_evaluated=len(df),
+                rows_failed=len(df),
+            )
+        duplicates = int(df[self.column].duplicated().sum())
+        status = CheckStatus.PASSED if duplicates == 0 else CheckStatus.FAILED
+        return CheckResult(
+            expectation_name="ExpectColumnValuesToBeUnique",
+            column=self.column,
+            status=status,
+            message=(
+                f"Column '{self.column}' has {duplicates} duplicate value(s)."
+                if duplicates > 0
+                else f"All values in '{self.column}' are unique."
+            ),
+            rows_evaluated=len(df),
+            rows_failed=duplicates,
+        )
+
+    def validate_spark(self, df: Any) -> CheckResult:
+        """Check uniqueness using PySpark."""
+        from pyspark.sql import functions as F  # noqa: N812, F401
+
+        total = df.count()
+        unique_count = df.select(self.column).distinct().count()
+        duplicates = total - unique_count
+        status = CheckStatus.PASSED if duplicates == 0 else CheckStatus.FAILED
+        return CheckResult(
+            expectation_name="ExpectColumnValuesToBeUnique",
+            column=self.column,
+            status=status,
+            message=(
+                f"Column '{self.column}' has {duplicates} duplicate value(s)."
+                if duplicates > 0
+                else f"All values in '{self.column}' are unique."
+            ),
+            rows_evaluated=total,
+            rows_failed=duplicates,
+        )
+
+
+@dataclass
+class ExpectColumnValuesToMatchRegex(BaseExpectation):
+    """Expect that all values in a column match a regular expression.
+
+    Parameters
+    ----------
+    column:
+        Name of the column to check.
+    pattern:
+        Regular expression pattern (Python :mod:`re` syntax).
+    """
+
+    column: str
+    pattern: str
+
+    def validate_pandas(self, df: Any) -> CheckResult:
+        """Check regex match using pandas."""
+        if self.column not in df.columns:
+            return CheckResult(
+                expectation_name="ExpectColumnValuesToMatchRegex",
+                column=self.column,
+                status=CheckStatus.FAILED,
+                message=f"Column '{self.column}' not found in DataFrame.",
+                rows_evaluated=len(df),
+                rows_failed=len(df),
+            )
+        series = df[self.column].astype(str)
+        non_matching = int((~series.str.match(self.pattern)).sum())
+        status = CheckStatus.PASSED if non_matching == 0 else CheckStatus.FAILED
+        return CheckResult(
+            expectation_name="ExpectColumnValuesToMatchRegex",
+            column=self.column,
+            status=status,
+            message=(
+                f"Column '{self.column}' has {non_matching} value(s) not matching "
+                f"pattern '{self.pattern}'."
+                if non_matching > 0
+                else f"All values in '{self.column}' match the pattern."
+            ),
+            rows_evaluated=len(df),
+            rows_failed=non_matching,
+            metadata={"pattern": self.pattern},
+        )
+
+    def validate_spark(self, df: Any) -> CheckResult:
+        """Check regex match using PySpark."""
+        from pyspark.sql import functions as F  # noqa: N812
+
+        total = df.count()
+        non_matching = df.filter(
+            ~F.col(self.column).cast("string").rlike(self.pattern)
+        ).count()
+        status = CheckStatus.PASSED if non_matching == 0 else CheckStatus.FAILED
+        return CheckResult(
+            expectation_name="ExpectColumnValuesToMatchRegex",
+            column=self.column,
+            status=status,
+            message=(
+                f"Column '{self.column}' has {non_matching} value(s) not matching "
+                f"pattern '{self.pattern}'."
+                if non_matching > 0
+                else f"All values in '{self.column}' match the pattern."
+            ),
+            rows_evaluated=total,
+            rows_failed=non_matching,
+            metadata={"pattern": self.pattern},
+        )
+
+
+__all__ = [
+    "BaseExpectation",
+    "ExpectColumnToNotBeNull",
+    "ExpectColumnValuesInRange",
+    "ExpectColumnValuesToBeUnique",
+    "ExpectColumnValuesToMatchRegex",
+]

aptdata/plugins/quality/report.py

@@ -0,0 +1,94 @@
+"""Quality check result and report types."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any
+
+
+class CheckStatus(str, Enum):
+    """Status of a single quality check."""
+
+    PASSED = "PASSED"
+    FAILED = "FAILED"
+    WARNING = "WARNING"
+
+
+@dataclass
+class CheckResult:
+    """Result of a single data quality expectation.
+
+    Parameters
+    ----------
+    expectation_name:
+        Class name or human-readable label of the expectation.
+    column:
+        Column name the check was applied to (empty for table-level checks).
+    status:
+        Whether the check passed, failed, or issued a warning.
+    message:
+        Human-readable description of the outcome.
+    rows_evaluated:
+        Total number of rows considered by the check.
+    rows_failed:
+        Number of rows that violated the expectation.
+    metadata:
+        Arbitrary extra metadata from the expectation.
+    """
+
+    expectation_name: str
+    column: str = ""
+    status: CheckStatus = CheckStatus.PASSED
+    message: str = ""
+    rows_evaluated: int = 0
+    rows_failed: int = 0
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class QualityReport:
+    """Aggregated quality report for a single dataset validation run.
+
+    Parameters
+    ----------
+    dataset_uri:
+        URI of the dataset that was validated.
+    workflow_name:
+        Name of the workflow that triggered validation.
+    trace_id:
+        OpenTelemetry trace identifier.
+    timestamp:
+        UTC ISO-8601 timestamp of when the report was generated.
+    checks:
+        Individual :class:`CheckResult` objects.
+    """
+
+    dataset_uri: str
+    workflow_name: str = ""
+    trace_id: str = ""
+    timestamp: str = field(
+        default_factory=lambda: datetime.now(timezone.utc).isoformat()
+    )
+    checks: list[CheckResult] = field(default_factory=list)
+
+    @property
+    def passed(self) -> bool:
+        """Return ``True`` when no check has a FAILED status."""
+        return all(c.status != CheckStatus.FAILED for c in self.checks)
+
+    @property
+    def summary(self) -> dict[str, int]:
+        """Return counts of PASSED, FAILED, and WARNING results."""
+        counts: dict[str, int] = {
+            CheckStatus.PASSED: 0,
+            CheckStatus.FAILED: 0,
+            CheckStatus.WARNING: 0,
+        }
+        for check in self.checks:
+            counts[check.status] += 1
+        return counts
+
+
+__all__ = ["CheckStatus", "CheckResult", "QualityReport"]