collate-data-diff 0.11.2__py3-none-any.whl

data_diff/abcs/database_types.py

@@ -0,0 +1,308 @@
+import decimal
+from abc import ABC, abstractmethod
+from typing import Collection, List, Optional, Tuple, Type, TypeVar, Union
+from datetime import datetime
+
+import attrs
+
+from data_diff.utils import ArithAlphanumeric, ArithUUID, Unknown
+
+
+DbPath = Tuple[str, ...]
+DbKey = Union[int, str, bytes, ArithUUID, ArithAlphanumeric]
+DbTime = datetime
+
+N = TypeVar("N")
+
+
+@attrs.frozen(kw_only=True, eq=False, order=False, unsafe_hash=True)
+class Collation:
+    """
+    A pre-parsed or pre-known record about db collation, per column.
+
+    The "greater" collation should be used as a target collation for textual PKs
+    on both sides of the diff — by coverting the "lesser" collation to self.
+
+    Snowflake easily absorbs the performance losses, so it has a boost to always
+    be greater than any other collation in non-Snowflake databases.
+    Other databases need to negotiate which side absorbs the performance impact.
+    """
+
+    # A boost for special databases that are known to absorb the performance dmaage well.
+    absorbs_damage: bool = False
+
+    # Ordinal soring by ASCII/UTF8 (True), or alphabetic as per locale/country/etc (False).
+    ordinal: Optional[bool] = None
+
+    # Lowercase first (aAbBcC or abcABC). Otherwise, uppercase first (AaBbCc or ABCabc).
+    lower_first: Optional[bool] = None
+
+    # 2-letter lower-case locale and upper-case country codes, e.g. en_US. Ignored for ordinals.
+    language: Optional[str] = None
+    country: Optional[str] = None
+
+    # There are also space-, punctuation-, width-, kana-(in)sensitivity, so on.
+    # Ignore everything not related to xdb alignment. Only case- & accent-sensitivity are common.
+    case_sensitive: Optional[bool] = None
+    accent_sensitive: Optional[bool] = None
+
+    # Purely informational, for debugging:
+    _source: Union[None, str, Collection[str]] = None
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Collation):
+            return NotImplemented
+        if self.ordinal and other.ordinal:
+            # TODO: does it depend on language? what does Albanic_BIN mean in MS SQL?
+            return True
+        return (
+            self.language == other.language
+            and (self.country is None or other.country is None or self.country == other.country)
+            and self.case_sensitive == other.case_sensitive
+            and self.accent_sensitive == other.accent_sensitive
+            and self.lower_first == other.lower_first
+        )
+
+    def __ne__(self, other: object) -> bool:
+        if not isinstance(other, Collation):
+            return NotImplemented
+        return not self.__eq__(other)
+
+    def __gt__(self, other: object) -> bool:
+        if not isinstance(other, Collation):
+            return NotImplemented
+        if self == other:
+            return False
+        if self.absorbs_damage and not other.absorbs_damage:
+            return False
+        if other.absorbs_damage and not self.absorbs_damage:
+            return True  # this one is preferred if it cannot absorb damage as its counterpart can
+        if self.ordinal and not other.ordinal:
+            return True
+        if other.ordinal and not self.ordinal:
+            return False
+        # TODO: try to align the languages & countries?
+        return False
+
+    def __ge__(self, other: object) -> bool:
+        if not isinstance(other, Collation):
+            return NotImplemented
+        return self == other or self.__gt__(other)
+
+    def __lt__(self, other: object) -> bool:
+        if not isinstance(other, Collation):
+            return NotImplemented
+        return self != other and not self.__gt__(other)
+
+    def __le__(self, other: object) -> bool:
+        if not isinstance(other, Collation):
+            return NotImplemented
+        return self == other or not self.__gt__(other)
+
+
+@attrs.define(frozen=True, kw_only=True)
+class ColType:
+    # Arbitrary metadata added and fetched at runtime.
+    _notes: List[N] = attrs.field(factory=list, init=False, hash=False, eq=False)
+
+    def add_note(self, note: N) -> None:
+        self._notes.append(note)
+
+    def get_note(self, cls: Type[N]) -> Optional[N]:
+        """Get the latest added note of type ``cls`` or its descendants."""
+        for note in reversed(self._notes):
+            if isinstance(note, cls):
+                return note
+        return None
+
+    @property
+    def supported(self) -> bool:
+        return True
+
+
+@attrs.define(frozen=True)
+class PrecisionType(ColType):
+    precision: int
+    rounds: Union[bool, Unknown] = Unknown
+
+
+@attrs.define(frozen=True)
+class Boolean(ColType):
+    precision = 0
+
+
+@attrs.define(frozen=True)
+class TemporalType(PrecisionType):
+    pass
+
+
+@attrs.define(frozen=True)
+class Timestamp(TemporalType):
+    pass
+
+
+@attrs.define(frozen=True)
+class TimestampTZ(TemporalType):
+    pass
+
+
+@attrs.define(frozen=True)
+class Datetime(TemporalType):
+    pass
+
+
+@attrs.define(frozen=True)
+class Date(TemporalType):
+    pass
+
+
+@attrs.define(frozen=True)
+class Time(TemporalType):
+    pass
+
+
+@attrs.define(frozen=True)
+class NumericType(ColType):
+    # 'precision' signifies how many fractional digits (after the dot) we want to compare
+    precision: int
+
+
+@attrs.define(frozen=True)
+class FractionalType(NumericType):
+    pass
+
+
+@attrs.define(frozen=True)
+class Float(FractionalType):
+    python_type = float
+
+
+@attrs.define(frozen=True)
+class IKey(ABC):
+    "Interface for ColType, for using a column as a key in table."
+
+    @property
+    @abstractmethod
+    def python_type(self) -> type:
+        "Return the equivalent Python type of the key"
+
+    def make_value(self, value):
+        if isinstance(value, self.python_type):
+            return value
+        return self.python_type(value)
+
+
+@attrs.define(frozen=True)
+class Decimal(FractionalType, IKey):  # Snowflake may use Decimal as a key
+    @property
+    def python_type(self) -> type:
+        if self.precision == 0:
+            return int
+        return decimal.Decimal
+
+
+@attrs.define(frozen=True)
+class StringType(ColType):
+    python_type = str
+    collation: Optional[Collation] = attrs.field(default=None, kw_only=True)
+
+
+@attrs.define(frozen=True)
+class ColType_UUID(ColType, IKey):
+    python_type = ArithUUID
+
+
+@attrs.define(frozen=True)
+class ColType_Alphanum(ColType, IKey):
+    python_type = ArithAlphanumeric
+
+
+@attrs.define(frozen=True)
+class Native_UUID(ColType_UUID):
+    pass
+
+
+@attrs.define(frozen=True)
+class String_UUID(ColType_UUID, StringType):
+    # Case is important for UUIDs stored as regular string, not native UUIDs stored as numbers.
+    # We slice them internally as numbers, but render them back to SQL as lower/upper case.
+    # None means we do not know for sure, behave as with False, but it might be unreliable.
+    lowercase: Optional[bool] = None
+    uppercase: Optional[bool] = None
+
+    def make_value(self, v: str) -> ArithUUID:
+        return self.python_type(v, lowercase=self.lowercase, uppercase=self.uppercase)
+
+
+@attrs.define(frozen=True)
+class String_Alphanum(ColType_Alphanum, StringType):
+    @staticmethod
+    def test_value(value: str) -> bool:
+        try:
+            ArithAlphanumeric(value)
+            return True
+        except ValueError:
+            return False
+
+
+@attrs.define(frozen=True)
+class String_VaryingAlphanum(String_Alphanum):
+    pass
+
+
+@attrs.define(frozen=True)
+class String_FixedAlphanum(String_Alphanum):
+    length: int
+
+    def make_value(self, value):
+        if isinstance(value, self.python_type):
+            return value
+        if len(value) != self.length:
+            raise ValueError(f"Expected alphanumeric value of length {self.length}, but got '{value}'.")
+        return self.python_type(value, max_len=self.length)
+
+
+@attrs.define(frozen=True)
+class Text(StringType):
+    @property
+    def supported(self) -> bool:
+        return False
+
+
+# In majority of DBMSes, it is called JSON/JSONB. Only in Snowflake, it is OBJECT.
+@attrs.define(frozen=True)
+class JSON(ColType):
+    pass
+
+
+@attrs.define(frozen=True)
+class Array(ColType):
+    item_type: ColType
+
+
+# Unlike JSON, structs are not free-form and have a very specific set of fields and their types.
+# We do not parse & use those fields now, but we can do this later.
+# For example, in BigQuery:
+# - https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#struct_type
+# - https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical#struct_literals
+@attrs.define(frozen=True)
+class Struct(ColType):
+    pass
+
+
+@attrs.define(frozen=True)
+class Integer(NumericType, IKey):
+    precision: int = 0
+    python_type: type = int
+
+    def __attrs_post_init__(self) -> None:
+        assert self.precision == 0
+
+
+@attrs.define(frozen=True)
+class UnknownColType(ColType):
+    text: str
+
+    @property
+    def supported(self) -> bool:
+        return False

data_diff/cloud/__init__.py

@@ -0,0 +1,2 @@
+from data_diff.cloud.datafold_api import DatafoldAPI, TCloudApiDataDiff, TCloudApiOrgMeta
+from data_diff.cloud.data_source import get_or_create_data_source

data_diff/cloud/data_source.py

@@ -0,0 +1,318 @@
+import json
+import time
+from typing import List, Optional, Union, overload
+
+import pydantic
+import rich
+from rich.table import Table
+from rich.prompt import Confirm, Prompt, FloatPrompt, IntPrompt, InvalidResponse
+from typing_extensions import Literal
+
+from data_diff.cloud.datafold_api import (
+    DatafoldAPI,
+    TCloudApiDataSourceConfigSchema,
+    TCloudApiDataSource,
+    TDsConfig,
+    TestDataSourceStatus,
+)
+from data_diff.dbt_parser import DbtParser
+
+
+UNKNOWN_VALUE = "unknown_value"
+
+
+class TDataSourceTestStage(pydantic.BaseModel):
+    name: str
+    status: TestDataSourceStatus
+    description: str = ""
+
+
+class TemporarySchemaPrompt(Prompt):
+    response_type = str
+
+    def process_response(self, value: str) -> str:
+        """Convert choices to a bool."""
+
+        if len(value.split(".")) != 2:
+            raise InvalidResponse("Temporary schema should have a format <database>.<schema>")
+        return value
+
+
+class ValueRequiredPrompt(Prompt):
+    def process_response(self, value: str) -> str:
+        value = super().process_response(value)
+        if value == UNKNOWN_VALUE or value is None or value == "":
+            raise InvalidResponse("Parameter must not be empty")
+        return value
+
+
+def _validate_temp_schema(temp_schema: str) -> None:
+    if len(temp_schema.split(".")) != 2:
+        raise ValueError("Temporary schema should have a format <database>.<schema>")
+
+
+def _get_temp_schema(dbt_parser: DbtParser, db_type: str) -> Optional[str]:
+    config = dbt_parser.get_datadiff_config()
+    config_prod_database = config.prod_database
+    config_prod_schema = config.prod_schema
+    if config_prod_database is not None and config_prod_schema is not None:
+        temp_schema = f"{config_prod_database}.{config_prod_schema}"
+        if db_type == "snowflake":
+            return temp_schema.upper()
+        elif db_type in {"pg", "postgres_aurora", "postgres_aws_rds", "redshift"}:
+            return temp_schema.lower()
+        return temp_schema
+    return
+
+
+def create_ds_config(
+    ds_config: TCloudApiDataSourceConfigSchema,
+    data_source_name: str,
+    dbt_parser: Optional[DbtParser] = None,
+) -> TDsConfig:
+    options = _parse_ds_credentials(ds_config=ds_config, only_basic_settings=True, dbt_parser=dbt_parser)
+
+    temp_schema = _get_temp_schema(dbt_parser=dbt_parser, db_type=ds_config.db_type) if dbt_parser else None
+    if temp_schema:
+        temp_schema = TemporarySchemaPrompt.ask("Temporary schema", default=temp_schema)
+    else:
+        temp_schema = TemporarySchemaPrompt.ask("Temporary schema (<database>.<schema>)")
+
+    float_tolerance = FloatPrompt.ask("Float tolerance", default=0.000001)
+
+    return TDsConfig(
+        name=data_source_name,
+        type=ds_config.db_type,
+        temp_schema=temp_schema,
+        float_tolerance=float_tolerance,
+        options=options,
+    )
+
+
+@overload
+def _cast_value(value: str, type_: Literal["integer"]) -> int: ...
+
+
+@overload
+def _cast_value(value: str, type_: Literal["boolean"]) -> bool: ...
+
+
+@overload
+def _cast_value(value: str, type_: Literal["string"]) -> str: ...
+
+
+def _cast_value(value: str, type_: str) -> Union[bool, int, str]:
+    if type_ == "integer":
+        return int(value)
+    elif type_ == "boolean":
+        return bool(value)
+    return value
+
+
+def _get_data_from_bigquery_json(path: str):
+    with open(path, "r") as file:
+        return json.load(file)
+
+
+def _align_dbt_cred_params_with_datafold_params(dbt_creds: dict) -> dict:
+    db_type = dbt_creds["type"]
+    if db_type == "bigquery":
+        method = dbt_creds["method"]
+        if method == "service-account":
+            data = _get_data_from_bigquery_json(path=dbt_creds["keyfile"])
+            dbt_creds["jsonKeyFile"] = json.dumps(data)
+        elif method == "service-account-json":
+            dbt_creds["jsonKeyFile"] = json.dumps(dbt_creds["keyfile_json"])
+        else:
+            rich.print(
+                f'[red]Cannot extract bigquery credentials from dbt_project.yml for "{method}" type. '
+                f"If you want to provide credentials via dbt_project.yml, "
+                f'please, use "service-account" or "service-account-json" '
+                f"(more in docs: https://docs.getdbt.com/reference/warehouse-setups/bigquery-setup). "
+                f"Otherwise, you can provide a path to a json key file or a json key file data as an input."
+            )
+        dbt_creds["projectId"] = dbt_creds["project"]
+    elif db_type == "snowflake":
+        dbt_creds["default_db"] = dbt_creds["database"]
+    elif db_type == "databricks":
+        dbt_creds["http_password"] = dbt_creds["token"]
+        dbt_creds["database"] = dbt_creds.get("catalog")
+    return dbt_creds
+
+
+def _parse_ds_credentials(
+    ds_config: TCloudApiDataSourceConfigSchema, only_basic_settings: bool = True, dbt_parser: Optional[DbtParser] = None
+):
+    creds = {}
+    use_dbt_data = False
+    if dbt_parser is not None:
+        use_dbt_data = Confirm.ask("Would you like to extract database credentials from dbt profiles.yml?")
+        try:
+            creds = dbt_parser.get_connection_creds()[0]
+            creds = _align_dbt_cred_params_with_datafold_params(dbt_creds=creds)
+        except Exception as e:
+            rich.print(f"[red]Cannot parse database credentials from dbt profiles.yml. Reason: {e}")
+
+    ds_options = {}
+    basic_required_fields = set(ds_config.config_schema.required)
+    for param_name, param_data in ds_config.config_schema.properties.items():
+        if only_basic_settings and param_name not in basic_required_fields:
+            continue
+
+        default_value = param_data.get("default", UNKNOWN_VALUE)
+        is_password = bool(param_data.get("format"))
+
+        title = param_data["title"]
+        type_ = param_data["type"]
+        input_values = {
+            "prompt": title,
+            "password": is_password,
+        }
+        if default_value != UNKNOWN_VALUE:
+            input_values["default"] = default_value
+
+        if use_dbt_data:
+            value = creds.get(param_name, UNKNOWN_VALUE)
+            if value == UNKNOWN_VALUE:
+                rich.print(f'[red]Cannot extract "{param_name}" from dbt profiles.yml. Please, type it manually')
+            else:
+                ds_options[param_name] = _cast_value(value, type_)
+                continue
+
+        if type_ == "integer":
+            value = IntPrompt.ask(**input_values)
+        elif type_ == "boolean":
+            value = Confirm.ask(title)
+        else:
+            value = ValueRequiredPrompt.ask(**input_values)
+
+        ds_options[param_name] = value
+    return ds_options
+
+
+def _check_data_source_exists(
+    data_sources: List[TCloudApiDataSource],
+    data_source_name: str,
+) -> Optional[TCloudApiDataSource]:
+    for ds in data_sources:
+        if ds.name == data_source_name:
+            return ds
+    return None
+
+
+def _test_data_source(api: DatafoldAPI, data_source_id: int, timeout: int = 64) -> List[TDataSourceTestStage]:
+    job_id = api.test_data_source(data_source_id)
+
+    checked_tests = {"connection", "temp_schema", "schema_download"}
+    seconds = 1
+    start = time.monotonic()
+    results = []
+    while True:
+        tests = api.check_data_source_test_results(job_id)
+        for test in tests:
+            if test.name not in checked_tests:
+                continue
+
+            if test.status == "done":
+                checked_tests.remove(test.name)
+                results.append(
+                    TDataSourceTestStage(name=test.name, status=test.result.status, description=test.result.message)
+                )
+
+        if not checked_tests:
+            break
+
+        if time.monotonic() - start > timeout:
+            for test_name in checked_tests:
+                results.append(
+                    TDataSourceTestStage(
+                        name=test_name,
+                        status=TestDataSourceStatus.SKIP,
+                        description=f"Does not complete in {timeout} seconds",
+                    )
+                )
+            break
+        time.sleep(seconds)
+        seconds *= 2
+
+    return results
+
+
+def _render_data_source(data_source: TCloudApiDataSource, title: str = "") -> None:
+    table = Table(title=title, min_width=80)
+    table.add_column("Parameter", justify="center", style="cyan")
+    table.add_column("Value", justify="center", style="magenta")
+    table.add_row("ID", str(data_source.id))
+    table.add_row("Name", data_source.name)
+    table.add_row("Type", data_source.type)
+    rich.print(table)
+
+
+def _render_available_data_sources(data_source_schema_configs: List[TCloudApiDataSourceConfigSchema]) -> None:
+    config_names = [ds_config.name for ds_config in data_source_schema_configs]
+
+    table = Table()
+    table.add_column("", justify="center", style="cyan")
+    table.add_column("Available data sources", style="magenta")
+    for i, db_type in enumerate(config_names, start=1):
+        table.add_row(str(i), db_type)
+    rich.print(table)
+
+
+def _render_data_source_test_results(test_results: List[TDataSourceTestStage]) -> None:
+    table = Table(title="Test results", min_width=80)
+    table.add_column(
+        "Test",
+        justify="center",
+        style="cyan",
+    )
+    table.add_column("Status", justify="center", style="magenta")
+    table.add_column("Description", justify="center", style="magenta")
+    for result in test_results:
+        table.add_row(result.name, result.status, result.description)
+    rich.print(table)
+
+
+def get_or_create_data_source(api: DatafoldAPI, dbt_parser: Optional[DbtParser] = None) -> int:
+    ds_configs = api.get_data_source_schema_config()
+    data_sources = api.get_data_sources()
+
+    _render_available_data_sources(data_source_schema_configs=ds_configs)
+    db_type_num = IntPrompt.ask(
+        prompt="What data source type do you want to create? Please, select a number",
+        choices=list(map(str, range(1, len(ds_configs) + 1))),
+        show_choices=False,
+    )
+
+    ds_config = ds_configs[db_type_num - 1]
+    default_ds_name = ds_config.name
+    rich.print("Press enter to accept the (Default value)")
+    ds_name = Prompt.ask("Data source name", default=default_ds_name)
+
+    ds = _check_data_source_exists(data_sources=data_sources, data_source_name=ds_name)
+    if ds is not None:
+        _render_data_source(data_source=ds, title=f'Found existing data source for name "{ds.name}"')
+        use_existing_ds = Confirm.ask("Would you like to continue with the existing data source?")
+        if not use_existing_ds:
+            return get_or_create_data_source(api=api, dbt_parser=dbt_parser)
+        return ds.id
+
+    ds_config = create_ds_config(ds_config=ds_config, data_source_name=ds_name, dbt_parser=dbt_parser)
+    ds = api.create_data_source(ds_config)
+    data_source_url = f"{api.host}/settings/integrations/dwh/{ds.type}/{ds.id}"
+    _render_data_source(data_source=ds, title=f"Created a new data source with ID = {ds.id} ({data_source_url})")
+
+    rich.print(
+        "We recommend to run tests for a new data source. "
+        "It requires some time but makes sure that the data source is configured correctly."
+    )
+    run_tests = Confirm.ask("Would you like to run tests?")
+    if run_tests:
+        test_results = _test_data_source(api=api, data_source_id=ds.id)
+        _render_data_source_test_results(test_results=test_results)
+        if any(result.status == TestDataSourceStatus.FAILED for result in test_results):
+            raise ValueError(
+                f"Data source tests failed. Please, try to update or test data source in the UI: {data_source_url}"
+            )
+
+    return ds.id