soda-bigquery 4.0.5__tar.gz

soda_bigquery-4.0.5/PKG-INFO

@@ -0,0 +1,6 @@
+Metadata-Version: 2.4
+Name: soda-bigquery
+Version: 4.0.5
+Requires-Dist: soda-core==4.0.5
+Requires-Dist: google-cloud-bigquery>=2.25.0
+Dynamic: requires-dist

soda_bigquery-4.0.5/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

soda_bigquery-4.0.5/setup.py

@@ -0,0 +1,24 @@
+#!/usr/bin/env python
+
+from setuptools import setup
+
+package_name = "soda-bigquery"
+package_version = "4.0.5"
+description = "Soda BigQuery V4"
+
+requires = [
+    f"soda-core=={package_version}",
+    "google-cloud-bigquery>=2.25.0",
+]
+
+setup(
+    name=package_name,
+    version=package_version,
+    install_requires=requires,
+    package_dir={"": "src"},
+    entry_points={
+        "soda.plugins.data_source.bigquery": [
+            "BigqueryDataSourceImpl = soda_bigquery.common.data_sources.bigquery_data_source:BigQueryDataSourceImpl",
+        ],
+    },
+)

soda_bigquery-4.0.5/src/soda_bigquery/common/data_sources/bigquery_data_source.py

@@ -0,0 +1,245 @@
+import logging
+from dataclasses import dataclass
+from typing import Optional
+
+from soda_bigquery.common.data_sources.bigquery_data_source_connection import (
+    BigQueryDataSource as BigQueryDataSourceModel,
+)
+from soda_bigquery.common.data_sources.bigquery_data_source_connection import (
+    BigQueryDataSourceConnection,
+)
+from soda_core.common.data_source_connection import DataSourceConnection
+from soda_core.common.data_source_impl import DataSourceImpl
+from soda_core.common.logging_constants import soda_logger
+from soda_core.common.metadata_types import DataSourceNamespace, SodaDataTypeName
+from soda_core.common.sql_ast import (
+    COLUMN,
+    CONCAT_WS,
+    COUNT,
+    DISTINCT,
+    LITERAL,
+    REGEX_LIKE,
+    STRING_HASH,
+    TUPLE,
+    VALUES,
+    WITH,
+)
+from soda_core.common.sql_dialect import SqlDialect
+from soda_core.common.statements.metadata_tables_query import MetadataTablesQuery
+
+logger: logging.Logger = soda_logger
+
+
+@dataclass
+class BigQueryDataSourceNamespace(DataSourceNamespace):
+    project_id: str
+    dataset: str
+    location: Optional[str] = None
+
+    def get_namespace_elements(self) -> list[str]:
+        if self.location:
+            return [self.location, self.project_id]
+        else:
+            return [self.project_id, self.dataset]
+
+    def get_database_for_metadata_query(self) -> str | None:
+        return self.project_id
+
+    def get_schema_for_metadata_query(self) -> str:
+        return self.dataset
+
+
+class BigQueryDataSourceImpl(DataSourceImpl, model_class=BigQueryDataSourceModel):
+    def __init__(self, data_source_model: BigQueryDataSourceModel, connection: Optional[DataSourceConnection] = None):
+        super().__init__(data_source_model=data_source_model, connection=connection)
+        self.cached_location = None
+
+    def _create_sql_dialect(self) -> SqlDialect:
+        return BigQuerySqlDialect(data_source_impl=self)
+
+    def _create_data_source_connection(self) -> DataSourceConnection:
+        return BigQueryDataSourceConnection(
+            name=self.data_source_model.name, connection_properties=self.data_source_model.connection_properties
+        )
+
+    def get_location(self) -> str:
+        if self.cached_location is not None:
+            location = self.cached_location
+        elif self.data_source_model.connection_properties.location is not None:
+            location = self.data_source_model.connection_properties.location
+        else:
+            result = self.execute_query("SELECT @@location")
+            location = result.rows[0][0]
+            logger.info(f"Detected BigQuery location: {location}")
+            self.cached_location = location
+        return location
+
+    def create_metadata_tables_query(self) -> MetadataTablesQuery:
+        super_metadata_tables_query = MetadataTablesQuery(
+            sql_dialect=self.sql_dialect,
+            data_source_connection=self.data_source_connection,
+            prefixes=[f"region-{self.get_location()}"],
+        )
+        return super_metadata_tables_query
+
+    def build_columns_metadata_query_str(self, dataset_prefixes: list[str], dataset_name: str) -> str:
+        table_namespace: DataSourceNamespace = BigQueryDataSourceNamespace(
+            project_id=dataset_prefixes[0], dataset=dataset_prefixes[1]
+        )
+
+        # BigQuery must be able to override to get the location
+        return self.sql_dialect.build_columns_metadata_query_str(
+            table_namespace=table_namespace, table_name=dataset_name
+        )
+
+    def _build_table_namespace_for_schema_query(self, prefixes: list[str]) -> tuple[DataSourceNamespace, str]:
+        table_namespace: DataSourceNamespace = BigQueryDataSourceNamespace(
+            project_id=prefixes[0],
+            dataset=None,  # We only need the project id to query the schemas, as it's always in the `INFORMATION_SCHEMA`
+        )
+        schema_name = self.extract_schema_from_prefix(prefixes)
+        if schema_name is None:
+            raise ValueError(f"Cannot determine schema name from prefixes: {prefixes}")
+
+        return table_namespace, schema_name
+
+
+class BigQuerySqlDialect(SqlDialect):
+    DEFAULT_QUOTE_CHAR = "`"
+
+    SODA_DATA_TYPE_SYNONYMS = (
+        (SodaDataTypeName.TEXT, SodaDataTypeName.VARCHAR, SodaDataTypeName.CHAR),
+        (SodaDataTypeName.INTEGER, SodaDataTypeName.BIGINT, SodaDataTypeName.SMALLINT),
+        (SodaDataTypeName.NUMERIC, SodaDataTypeName.DECIMAL),
+        (SodaDataTypeName.TIMESTAMP, SodaDataTypeName.TIMESTAMP_TZ),  # Bigquery does not support timezones
+        (SodaDataTypeName.FLOAT, SodaDataTypeName.DOUBLE),
+    )
+
+    def get_data_source_data_type_name_by_soda_data_type_names(self) -> dict[str, str]:
+        return {
+            SodaDataTypeName.CHAR: "string",  # BigQuery only has STRING
+            SodaDataTypeName.VARCHAR: "string",
+            SodaDataTypeName.TEXT: "string",  # alias for varchar
+            SodaDataTypeName.SMALLINT: "int64",  # all integers → INT64
+            SodaDataTypeName.INTEGER: "int64",
+            SodaDataTypeName.BIGINT: "int64",
+            SodaDataTypeName.DECIMAL: "numeric",  # NUMERIC (fixed precision: 38 digits, 9 decimals)
+            SodaDataTypeName.NUMERIC: "numeric",
+            SodaDataTypeName.FLOAT: "float64",  # FLOAT & DOUBLE → FLOAT64
+            SodaDataTypeName.DOUBLE: "float64",
+            SodaDataTypeName.TIMESTAMP: "timestamp",  # UTC-based
+            SodaDataTypeName.TIMESTAMP_TZ: "timestamp",  # still just TIMESTAMP in BigQuery
+            SodaDataTypeName.DATE: "date",
+            SodaDataTypeName.TIME: "time",
+            SodaDataTypeName.BOOLEAN: "bool",
+        }
+
+    def get_soda_data_type_name_by_data_source_data_type_names(self) -> dict[str, SodaDataTypeName]:
+        return {
+            "string": SodaDataTypeName.VARCHAR,
+            "smallint": SodaDataTypeName.SMALLINT,
+            "int64": SodaDataTypeName.BIGINT,
+            "numeric": SodaDataTypeName.NUMERIC,
+            "float64": SodaDataTypeName.DOUBLE,
+            "double precision": SodaDataTypeName.DOUBLE,
+            "timestamp": SodaDataTypeName.TIMESTAMP,
+            "date": SodaDataTypeName.DATE,
+            "time": SodaDataTypeName.TIME,
+            "bool": SodaDataTypeName.BOOLEAN,
+        }
+
+    def _get_data_type_name_synonyms(self) -> list[list[str]]:
+        return [
+            ["int64", "integer"],
+            ["bool", "boolean"],
+            ["numeric", "decimal"],
+            ["bignumeric", "bigdecimal"],
+        ]
+
+    def information_schema_namespace_elements(self, data_source_namespace: BigQueryDataSourceNamespace) -> list[str]:
+        return [data_source_namespace.project_id, data_source_namespace.dataset, "INFORMATION_SCHEMA"]
+
+    def default_casify(self, identifier: str) -> str:
+        return identifier.upper()
+
+    def _build_tuple_sql(self, tuple: TUPLE) -> str:
+        if tuple.check_context(COUNT) and tuple.check_context(DISTINCT):
+            return self._build_tuple_sql_in_distinct(tuple)
+        return f"{super()._build_tuple_sql(tuple)}"
+
+    def _build_tuple_sql_in_distinct(self, tuple: TUPLE) -> str:
+        return f"TO_JSON_STRING(STRUCT({super()._build_tuple_sql(tuple)}))"
+
+    def _build_regex_like_sql(self, matches: REGEX_LIKE) -> str:
+        expression: str = self.build_expression_sql(matches.expression)
+        return f"REGEXP_CONTAINS({expression}, r'{matches.regex_pattern}')"
+
+    def supports_data_type_character_maximum_length(self) -> bool:
+        return False
+
+    def supports_data_type_numeric_precision(self) -> bool:
+        return False
+
+    def supports_cte_alias_columns(self) -> bool:
+        return False
+
+    def supports_data_type_numeric_scale(self) -> bool:
+        return False
+
+    def supports_data_type_datetime_precision(self) -> bool:
+        return False
+
+    def sql_expr_timestamp_literal(self, datetime_in_iso8601: str) -> str:
+        return f"timestamp('{datetime_in_iso8601}')"
+
+    def sql_expr_timestamp_truncate_day(self, timestamp_literal: str) -> str:
+        return f"date_trunc(timestamp({timestamp_literal}), day)"
+
+    def sql_expr_timestamp_add_day(self, timestamp_literal: str) -> str:
+        return f"{timestamp_literal} + interval 1 day"
+
+    def literal_string(self, value: str) -> Optional[str]:
+        if value is None:
+            return None
+        return "'''" + self.escape_string(value) + "'''"
+
+    def build_cte_values_sql(self, values: VALUES, alias_columns: list[COLUMN] | None) -> str:
+        # The first select row should have column aliases
+        # Remaining rows don't need aliases
+        def build_literal_with_alias(literal, alias: COLUMN | None) -> str:
+            return f"{self.literal(literal)} AS {self.quote_column(alias.name)}"
+
+        literal_rows: list[str] = []
+        for tuple in values.values:
+            if alias_columns:
+                literal_sqls: list[str] = []
+                for i in range(len(tuple.expressions)):
+                    literal: LITERAL = tuple.expressions[i]
+                    alias: COLUMN = alias_columns[i]
+                    literal_sqls.append(build_literal_with_alias(literal=literal, alias=alias))
+                literal_rows.append(", ".join(literal_sql for literal_sql in literal_sqls))
+                alias_columns = None
+            else:
+                literal_rows.append(", ".join(self.build_expression_sql(e) for e in tuple.expressions))
+
+        select_rows: list[str] = [f"SELECT {literal_row}" for literal_row in literal_rows]
+
+        return "\nUNION ALL ".join([select_row for select_row in select_rows])
+
+    def _build_cte_with_sql_line(self, with_element: WITH) -> str:
+        return f"WITH {self.quote_default(with_element.alias)} AS ("
+
+    def get_max_sql_statement_length(self) -> int:
+        # This is a VERY conservative limit. We can probably get away with more
+        # However, BigQuery started complaining that the statement was "too complicated" with a length of 1MB.
+        return 200 * 1024  # 200KB
+
+    def get_preferred_number_of_rows_for_insert(self) -> int:
+        return 500
+
+    def _build_concat_ws_sql(self, concat_ws: CONCAT_WS) -> str:
+        elements: str = f", '{concat_ws.separator}', ".join(self.build_expression_sql(e) for e in concat_ws.expressions)
+        return f"CONCAT({elements})"
+
+    def _build_string_hash_sql(self, string_hash: STRING_HASH) -> str:
+        return f"to_hex({super()._build_string_hash_sql(string_hash)})"

soda_bigquery-4.0.5/src/soda_bigquery/common/data_sources/bigquery_data_source_connection.py

@@ -0,0 +1,152 @@
+from __future__ import annotations
+
+import json
+import logging
+from abc import ABC
+from typing import Literal, Optional, Union
+
+from google.api_core.client_info import ClientInfo
+from google.auth import default, impersonated_credentials
+from google.cloud import bigquery
+from google.cloud.bigquery import dbapi
+from google.cloud.bigquery.table import Row
+from google.oauth2.service_account import Credentials
+from pydantic import Field, SecretStr
+from soda_core.common.data_source_connection import DataSourceConnection
+from soda_core.common.logging_constants import soda_logger
+from soda_core.model.data_source.data_source import DataSourceBase
+from soda_core.model.data_source.data_source_connection_properties import (
+    DataSourceConnectionProperties,
+)
+
+logger: logging.Logger = soda_logger
+
+
+CONTEXT_AUTHENTICATION_DESCRIPTION = "Use context authentication"
+
+
+class BigQueryConnectionProperties(DataSourceConnectionProperties, ABC):
+    project_id: Optional[str] = Field(None, description="BigQuery project ID")
+    storage_project_id: Optional[str] = Field(None, description="BigQuery storage project ID")
+    location: Optional[str] = Field(None, description="BigQuery location")
+    client_options: Optional[dict] = Field(None, description="Client options")
+    labels: Optional[dict] = Field({}, description="Labels")
+    auth_scopes: Optional[list[str]] = Field(
+        [
+            "https://www.googleapis.com/auth/bigquery",
+            "https://www.googleapis.com/auth/cloud-platform",
+            "https://www.googleapis.com/auth/drive",
+        ],
+        description="Authentication scopes",
+    )
+    impersonation_account: Optional[str] = Field(None, description="Impersonation account")
+    delegates: Optional[list[str]] = Field(None, description="Delegates")
+
+
+class BigQueryJSONStringAuth(BigQueryConnectionProperties):
+    """BigQuery authentication using JSON string"""
+
+    use_context_auth: Optional[Literal[False]] = Field(False, description=CONTEXT_AUTHENTICATION_DESCRIPTION)
+    account_info_json: SecretStr = Field(..., description="Service account JSON as string", min_length=1)
+
+
+class BigQueryJSONFileAuth(BigQueryConnectionProperties):
+    """BigQuery authentication using JSON file path"""
+
+    use_context_auth: Optional[Literal[False]] = Field(False, description=CONTEXT_AUTHENTICATION_DESCRIPTION)
+    account_info_json_path: str = Field(..., description="Path to service account JSON file", min_length=1)
+
+
+class BigQueryContextAuth(BigQueryConnectionProperties):
+    """BigQuery authentication using context.
+
+    If use_context_auth is True, then application default credentials will be used.
+    The user may optionally provide JSON credentials; they will be ignored.
+    """
+
+    use_context_auth: Literal[True] = Field(description=CONTEXT_AUTHENTICATION_DESCRIPTION)
+
+
+class BigQueryDataSource(DataSourceBase, ABC):
+    type: Literal["bigquery"] = Field("bigquery")
+
+    connection_properties: Union[
+        BigQueryJSONStringAuth,
+        BigQueryJSONFileAuth,
+        BigQueryContextAuth,
+    ] = Field(..., alias="connection", description="BigQuery connection configuration")
+
+
+class BigQueryDataSourceConnection(DataSourceConnection):
+    def __init__(self, name: str, connection_properties: DataSourceConnectionProperties):
+        super().__init__(name, connection_properties)
+
+    def _load_project_id_and_credentials(self, config: BigQueryConnectionProperties):
+        if isinstance(config, BigQueryContextAuth):
+            logger.info("Using application default credentials.")
+            self.credentials, self.project_id = default()
+            return
+
+        if isinstance(config, BigQueryJSONFileAuth):
+            account_info_dict = json.load(open(config.account_info_json_path))
+        elif isinstance(config, BigQueryJSONStringAuth):
+            account_info_dict = json.loads(config.account_info_json.get_secret_value())
+        self.credentials = Credentials.from_service_account_info(
+            account_info_dict,
+            scopes=config.auth_scopes,
+        )
+        self.project_id = account_info_dict.get("project_id")
+
+    def _load_optional_impersonated_credentials(self, config: BigQueryConnectionProperties):
+        if config.impersonation_account:
+            logger.info("Using impersonation of Service Account.")
+            if config.delegates:
+                logger.info("Using Service Account delegates.")
+                delegates = config.delegates
+            else:
+                delegates = None
+            self.credentials = impersonated_credentials.Credentials(
+                source_credentials=self.credentials,
+                target_principal=str(config.impersonation_account),
+                target_scopes=config.auth_scopes,
+                delegates=delegates,
+            )
+
+    def _apply_optional_params(self, config: BigQueryConnectionProperties):
+        # Users can optionally overwrite in the connection properties
+        self.project_id = config.project_id if config.project_id else self.project_id
+        self.location = config.location
+        self.client_options = config.client_options
+
+        # Storage project ID is currently not used, because the project is configured via the DQN in the data contract.
+        # When we implement discovery, we'll need to use this value.
+        self.storage_project_id = config.storage_project_id if config.storage_project_id else self.project_id
+
+        self.labels = config.labels
+
+    def _create_connection(
+        self,
+        config: BigQueryConnectionProperties,
+    ):
+        self._load_project_id_and_credentials(config)
+        self._load_optional_impersonated_credentials(config)
+        self._apply_optional_params(config)
+
+        client_info = ClientInfo(
+            user_agent="soda-core",
+        )
+        default_query_job_config = bigquery.QueryJobConfig(labels=self.labels)
+        self.client = bigquery.Client(
+            project=self.project_id,
+            credentials=self.credentials,
+            default_query_job_config=default_query_job_config,
+            client_info=client_info,
+            location=config.location,
+            client_options=self.client_options,
+        )
+
+        return dbapi.Connection(self.client)
+
+    def format_rows(self, rows: list[Row]) -> list[tuple]:
+        formatted_rows = [tuple(r.values()) for r in rows]
+        return formatted_rows

soda_bigquery-4.0.5/src/soda_bigquery/test_helpers/bigquery_data_source_test_helper.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import json
+import os
+
+from helpers.data_source_test_helper import DataSourceTestHelper
+
+
+class BigQueryDataSourceTestHelper(DataSourceTestHelper):
+    def _create_database_name(self) -> str | None:
+        # Parse the dataset name from the account info json
+        account_info_json = json.loads(os.getenv("BIGQUERY_ACCOUNT_INFO_JSON", "{}"))
+        database_name = account_info_json.get("project_id", "soda-testing-dataset")
+        return database_name
+
+    def _create_data_source_yaml_str(self) -> str:
+        """
+        Called in _create_data_source_impl to initialized self.data_source_impl
+        self.database_name and self.schema_name are available if appropriate for the data source type
+        """
+        return f"""
+            type: bigquery
+            name: {self.name}
+            connection:
+                account_info_json: '{os.getenv("BIGQUERY_ACCOUNT_INFO_JSON", "")}'
+
+        """
+        # location: '{os.getenv("BIGQUERY_LOCATION", "US")}'
+
+    def _get_contract_data_type_dict(self) -> dict[str, str]:
+        """
+        DataSourceTestHelpers can override this method as an easy way
+        to customize the get_schema_check_sql_type behavior
+        """
+        return self._get_create_table_sql_type_dict()

soda_bigquery-4.0.5/src/soda_bigquery.egg-info/PKG-INFO

@@ -0,0 +1,6 @@
+Metadata-Version: 2.4
+Name: soda-bigquery
+Version: 4.0.5
+Requires-Dist: soda-core==4.0.5
+Requires-Dist: google-cloud-bigquery>=2.25.0
+Dynamic: requires-dist

soda_bigquery-4.0.5/src/soda_bigquery.egg-info/SOURCES.txt

@@ -0,0 +1,10 @@
+setup.py
+src/soda_bigquery.egg-info/PKG-INFO
+src/soda_bigquery.egg-info/SOURCES.txt
+src/soda_bigquery.egg-info/dependency_links.txt
+src/soda_bigquery.egg-info/entry_points.txt
+src/soda_bigquery.egg-info/requires.txt
+src/soda_bigquery.egg-info/top_level.txt
+src/soda_bigquery/common/data_sources/bigquery_data_source.py
+src/soda_bigquery/common/data_sources/bigquery_data_source_connection.py
+src/soda_bigquery/test_helpers/bigquery_data_source_test_helper.py

soda_bigquery-4.0.5/src/soda_bigquery.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

soda_bigquery-4.0.5/src/soda_bigquery.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[soda.plugins.data_source.bigquery]
+BigqueryDataSourceImpl = soda_bigquery.common.data_sources.bigquery_data_source:BigQueryDataSourceImpl

soda_bigquery-4.0.5/src/soda_bigquery.egg-info/requires.txt

@@ -0,0 +1,2 @@
+soda-core==4.0.5
+google-cloud-bigquery>=2.25.0

soda_bigquery-4.0.5/src/soda_bigquery.egg-info/top_level.txt

@@ -0,0 +1 @@
+soda_bigquery