diff --git a/docsite/docs/connectors/bigquery.md b/docsite/docs/connectors/bigquery.md new file mode 100644 index 0000000..a0d7a44 --- /dev/null +++ b/docsite/docs/connectors/bigquery.md @@ -0,0 +1,109 @@ +--- +sidebar_position: 7 +--- + +# BigQuery + +`intugle` integrates with Google BigQuery, allowing you to read data from your datasets for profiling, analysis, and data product generation. + +## Installation + +To use `intugle` with BigQuery, you must install the optional dependencies: + +```bash +pip install "intugle[bigquery]" +``` + +This installs the `google-cloud-bigquery` library. + +## Configuration + +To connect to your BigQuery project, you must provide connection credentials in a `profiles.yml` file at the root of your project. The adapter looks for a top-level `bigquery:` key. + +**Example `profiles.yml`:** + +```yaml +bigquery: + name: my_bigquery_source + project_id: + dataset: + location: US # Optional, defaults to US + credentials_path: /path/to/service-account-credentials.json # Optional +``` + +### Authentication Options + +1. **Service Account JSON File** (Recommended for production): + - Set `credentials_path` to your service account JSON file. + - The service account needs **BigQuery Data Viewer** and **BigQuery Job User** roles. + +2. **Application Default Credentials** (For development): + - Omit `credentials_path`. + - Uses `gcloud auth application-default login`. + - Or uses environment variable `GOOGLE_APPLICATION_CREDENTIALS`. + +## Usage + +### Reading Data from BigQuery + +To include a BigQuery table or view in your `SemanticModel`, define it in your input dictionary with `type: "bigquery"` and use the `identifier` key to specify the table name. + +:::caution Important +The dictionary key for your dataset (e.g., `"my_table"`) must exactly match the table name specified in the `identifier`. +::: + +```python +from intugle import SemanticModel + +datasets = { + "my_table": { + "identifier": "my_table", # Must match the key above + "type": "bigquery" + }, + "another_view": { + "identifier": "another_view", + "type": "bigquery" + } +} + +# Initialize the semantic model +sm = SemanticModel(datasets, domain="Analytics") + +# Build the model as usual +sm.build() +``` + +### Materializing Data Products + +When you use the `DataProduct` class with a BigQuery connection, the resulting data product can be materialized as a new **table** or **view** directly within your target dataset. + +:::caution +**Beta Feature:** The DataProduct feature for BigQuery is currently in beta. If you encounter any issues, please raise them on our [GitHub issues page](https://github.com/Intugle/data-tools/issues). +::: + +```python +from intugle import DataProduct + +etl_model = { + "name": "top_users", + "fields": [ + {"id": "users.id", "name": "user_id"}, + {"id": "users.name", "name": "user_name"}, + ] +} + +dp = DataProduct() + +# Materialize as a view (default) +dp.build(etl_model, materialize="view") + +# Materialize as a table +dp.build(etl_model, materialize="table") +``` + +:::info Required Permissions +To successfully materialise data products, the Service Account or User must have the following privileges: +* `roles/bigquery.dataViewer` - Read table data +* `roles/bigquery.jobUser` - Run queries +* `roles/bigquery.dataEditor` - Create tables and views +::: diff --git a/docsite/docs/connectors/implementing-a-connector.md b/docsite/docs/connectors/implementing-a-connector.md index ff6b5d6..c7458e1 100644 --- a/docsite/docs/connectors/implementing-a-connector.md +++ b/docsite/docs/connectors/implementing-a-connector.md @@ -1,5 +1,5 @@ --- -sidebar_position: 7 +sidebar_position: 8 --- # Implementing a Connector diff --git a/pyproject.toml b/pyproject.toml index 4f0a535..21329e1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,6 +69,10 @@ postgres = [ "asyncpg>=0.30.0", "sqlglot>=27.20.0", ] +bigquery = [ + "google-cloud-bigquery>=3.11.0", + "sqlglot>=27.20.0", +] sqlserver = [ "mssql-python>=0.13.1", "sqlglot>=27.20.0", diff --git a/src/intugle/adapters/factory.py b/src/intugle/adapters/factory.py index 9421636..d3105df 100644 --- a/src/intugle/adapters/factory.py +++ b/src/intugle/adapters/factory.py @@ -36,6 +36,7 @@ def is_safe_plugin_name(plugin_name: str) -> bool: "intugle.adapters.types.mariadb.mariadb", "intugle.adapters.types.sqlserver.sqlserver", "intugle.adapters.types.sqlite.sqlite", + "intugle.adapters.types.bigquery.bigquery", "intugle.adapters.types.oracle.oracle", ] diff --git a/src/intugle/adapters/models.py b/src/intugle/adapters/models.py index 3f43338..97233f8 100644 --- a/src/intugle/adapters/models.py +++ b/src/intugle/adapters/models.py @@ -19,6 +19,7 @@ def get_dataset_data_type() -> type: if TYPE_CHECKING: import pandas as pd + from intugle.adapters.types.bigquery.models import BigQueryConfig from intugle.adapters.types.databricks.models import DatabricksConfig from intugle.adapters.types.duckdb.models import DuckdbConfig from intugle.adapters.types.mariadb.models import MariaDBConfig @@ -28,7 +29,7 @@ def get_dataset_data_type() -> type: from intugle.adapters.types.sqlite.models import SqliteConfig from intugle.adapters.types.sqlserver.models import SQLServerConfig - DataSetData = pd.DataFrame | DuckdbConfig | SnowflakeConfig | DatabricksConfig | PostgresConfig | SQLServerConfig | SqliteConfig | OracleConfig | MariaDBConfig + DataSetData = pd.DataFrame | DuckdbConfig | SnowflakeConfig | DatabricksConfig | PostgresConfig | SQLServerConfig | SqliteConfig | OracleConfig | MariaDBConfig | BigQueryConfig else: # At runtime, this is dynamically determined DataSetData = Any diff --git a/src/intugle/adapters/types/bigquery/__init__.py b/src/intugle/adapters/types/bigquery/__init__.py new file mode 100644 index 0000000..0547d74 --- /dev/null +++ b/src/intugle/adapters/types/bigquery/__init__.py @@ -0,0 +1 @@ +# BigQuery adapter for Intugle diff --git a/src/intugle/adapters/types/bigquery/bigquery.py b/src/intugle/adapters/types/bigquery/bigquery.py new file mode 100644 index 0000000..f121a05 --- /dev/null +++ b/src/intugle/adapters/types/bigquery/bigquery.py @@ -0,0 +1,407 @@ +import time +from typing import TYPE_CHECKING, Any, Optional + +import numpy as np +import pandas as pd + +from intugle.adapters.adapter import Adapter +from intugle.adapters.factory import AdapterFactory +from intugle.adapters.models import ColumnProfile, DataSetData, ProfilingOutput +from intugle.adapters.types.bigquery.models import BigQueryConfig, BigQueryConnectionConfig +from intugle.adapters.utils import convert_to_native +from intugle.core import settings +from intugle.core.utilities.processing import string_standardization + +if TYPE_CHECKING: + from intugle.analysis.models import DataSet + +try: + from google.cloud import bigquery + from google.oauth2 import service_account + from sqlglot import transpile + + BIGQUERY_AVAILABLE = True +except ImportError: + BIGQUERY_AVAILABLE = False + + +class BigQueryAdapter(Adapter): + _instance = None + _initialized = False + + @property + def database(self) -> Optional[str]: + return self._project_id + + @database.setter + def database(self, value: str): + self._project_id = value + + @property + def schema(self) -> Optional[str]: + return self._dataset_id + + @schema.setter + def schema(self, value: str): + self._dataset_id = value + + @property + def source_name(self) -> str: + return self._source_name + + @source_name.setter + def source_name(self, value: str): + self._source_name = value + + def __new__(cls, *args, **kwargs): + if not cls._instance: + cls._instance = super().__new__(cls) + return cls._instance + + def __init__(self): + if self._initialized: + return + + if not BIGQUERY_AVAILABLE: + raise ImportError( + "BigQuery dependencies are not installed. Please run 'pip install intugle[bigquery]'." + ) + + self.client: Optional["bigquery.Client"] = None + self._project_id: Optional[str] = None + self._dataset_id: Optional[str] = None + self._location: str = "US" + self._source_name: str = settings.PROFILES.get("bigquery", {}).get("name", "my_bigquery_source") + + self.connect() + self._initialized = True + + def connect(self): + """Establish connection to BigQuery.""" + connection_parameters_dict = settings.PROFILES.get("bigquery", {}) + if not connection_parameters_dict: + raise ValueError("Could not create BigQuery connection. No 'bigquery' section found in profiles.yml.") + + params = BigQueryConnectionConfig.model_validate(connection_parameters_dict) + self._project_id = params.project_id + self._dataset_id = params.dataset_id + self._location = params.location + + # Initialize BigQuery client with credentials if provided + if params.credentials_path: + credentials = service_account.Credentials.from_service_account_file(params.credentials_path) + self.client = bigquery.Client(project=self._project_id, credentials=credentials, location=self._location) + else: + # Use default credentials (Application Default Credentials) + self.client = bigquery.Client(project=self._project_id, location=self._location) + + def _get_fqn(self, identifier: str) -> str: + """Gets the fully qualified name for a table identifier.""" + if "." in identifier: + # Already has project or dataset prefix + parts = identifier.split(".") + if len(parts) == 2: + # dataset.table format + return f"`{self._project_id}.{identifier}`" + elif len(parts) == 3: + # project.dataset.table format + return f"`{identifier}`" + return f"`{self._project_id}.{self._dataset_id}.{identifier}`" + + @staticmethod + def check_data(data: Any) -> BigQueryConfig: + try: + data = BigQueryConfig.model_validate(data) + except Exception: + raise TypeError("Input must be a BigQuery config.") + return data + + def _execute_sql(self, query: str) -> list[Any]: + """Execute a SQL query and return results as a list of rows.""" + job_config = bigquery.QueryJobConfig(default_dataset=f"{self._project_id}.{self._dataset_id}") + query_job = self.client.query(query, job_config=job_config) + results = query_job.result() + return [dict(row) for row in results] + + def _get_pandas_df(self, query: str) -> pd.DataFrame: + """Execute a SQL query and return results as a pandas DataFrame.""" + job_config = bigquery.QueryJobConfig(default_dataset=f"{self._project_id}.{self._dataset_id}") + query_job = self.client.query(query, job_config=job_config) + return query_job.to_dataframe() + + def profile(self, data: BigQueryConfig, table_name: str) -> ProfilingOutput: + """Profile a BigQuery table.""" + data = self.check_data(data) + fqn = self._get_fqn(data.identifier) + + # Get total count + count_query = f"SELECT COUNT(*) as count FROM {fqn}" + total_count = self._execute_sql(count_query)[0]["count"] + + # Get column information from INFORMATION_SCHEMA + schema_query = f""" + SELECT column_name, data_type + FROM `{self._project_id}.{self._dataset_id}.INFORMATION_SCHEMA.COLUMNS` + WHERE table_name = @table_name + ORDER BY ordinal_position + """ + job_config = bigquery.QueryJobConfig( + query_parameters=[bigquery.ScalarQueryParameter("table_name", "STRING", data.identifier)] + ) + query_job = self.client.query(schema_query, job_config=job_config) + rows = [dict(row) for row in query_job.result()] + + columns = [row["column_name"] for row in rows] + dtypes = {row["column_name"]: row["data_type"] for row in rows} + + return ProfilingOutput( + count=total_count, + columns=columns, + dtypes=dtypes, + ) + + def column_profile( + self, + data: BigQueryConfig, + table_name: str, + column_name: str, + total_count: int, + sample_limit: int = 10, + dtype_sample_limit: int = 10000, + ) -> Optional[ColumnProfile]: + """Profile a specific column in a BigQuery table.""" + data = self.check_data(data) + fqn = self._get_fqn(data.identifier) + start_ts = time.time() + + # Null and distinct counts + query = f""" + SELECT + COUNTIF(`{column_name}` IS NULL) as null_count, + COUNT(DISTINCT `{column_name}`) as distinct_count + FROM {fqn} + """ + result = self._execute_sql(query)[0] + null_count = result["null_count"] + distinct_count = result["distinct_count"] + not_null_count = total_count - null_count + + # Sampling for distinct values + sample_query = f""" + SELECT DISTINCT CAST(`{column_name}` AS STRING) as value + FROM {fqn} + WHERE `{column_name}` IS NOT NULL + LIMIT {dtype_sample_limit} + """ + distinct_values_result = self._execute_sql(sample_query) + distinct_values = [row["value"] for row in distinct_values_result if row["value"] is not None] + + if distinct_count > 0 and len(distinct_values) > 0: + distinct_sample_size = min(distinct_count, dtype_sample_limit, len(distinct_values)) + sample_data = list(np.random.choice(distinct_values, distinct_sample_size, replace=False)) + else: + sample_data = [] + + dtype_sample = None + if distinct_count >= dtype_sample_limit: + dtype_sample = sample_data + elif distinct_count > 0 and not_null_count > 0: + remaining_sample_size = dtype_sample_limit - len(distinct_values) + if remaining_sample_size > 0: + additional_samples_query = f""" + SELECT CAST(`{column_name}` AS STRING) as value + FROM {fqn} + WHERE `{column_name}` IS NOT NULL + ORDER BY RAND() + LIMIT {remaining_sample_size} + """ + additional_samples_result = self._execute_sql(additional_samples_query) + additional_samples = [row["value"] for row in additional_samples_result if row["value"] is not None] + dtype_sample = list(distinct_values) + additional_samples + else: + dtype_sample = list(distinct_values) + else: + dtype_sample = [] + + native_sample_data = convert_to_native(sample_data) + native_dtype_sample = convert_to_native(dtype_sample) + business_name = string_standardization(column_name) + + return ColumnProfile( + column_name=column_name, + table_name=table_name, + business_name=business_name, + null_count=null_count, + count=total_count, + distinct_count=distinct_count, + uniqueness=distinct_count / total_count if total_count > 0 else 0.0, + completeness=not_null_count / total_count if total_count > 0 else 0.0, + sample_data=native_sample_data[:sample_limit], + dtype_sample=native_dtype_sample, + ts=time.time() - start_ts, + ) + + def load(self, data: BigQueryConfig, table_name: str): + """Load data into BigQuery table.""" + data = self.check_data(data) + # This method is typically used for loading data from other sources + # Implementation depends on the specific use case + # raise NotImplementedError("Load method needs to be implemented based on specific requirements.") + + def execute(self, query: str): + """Execute a SQL query.""" + job_config = bigquery.QueryJobConfig(default_dataset=f"{self._project_id}.{self._dataset_id}") + query_job = self.client.query(query, job_config=job_config) + query_job.result() # Wait for the query to complete + return query_job + + def to_df(self, data: DataSetData, table_name: str) -> pd.DataFrame: + """Convert BigQuery table to pandas DataFrame.""" + data = self.check_data(data) + fqn = self._get_fqn(data.identifier) + query = f"SELECT * FROM {fqn}" + return self._get_pandas_df(query) + + def to_df_from_query(self, query: str) -> pd.DataFrame: + """Execute a query and return results as DataFrame.""" + return self._get_pandas_df(query) + + def create_table_from_query( + self, table_name: str, query: str, materialize: str = "view", **kwargs + ) -> str: + """Create a table or view from a query.""" + fqn = self._get_fqn(table_name) + transpiled_sql = transpile(query, write="bigquery")[0] + + if materialize == "view": + create_query = f"CREATE OR REPLACE VIEW {fqn} AS {transpiled_sql}" + elif materialize == "table": + create_query = f"CREATE OR REPLACE TABLE {fqn} AS {transpiled_sql}" + else: + raise ValueError(f"Invalid materialize option: {materialize}. Use 'table' or 'view'.") + + self.execute(create_query) + return transpiled_sql + + def create_new_config_from_etl(self, etl_name: str) -> DataSetData: + """Create a new config for an ETL-generated table.""" + return BigQueryConfig(identifier=etl_name, type="bigquery") + + def intersect_count( + self, table1: "DataSet", column1_name: str, table2: "DataSet", column2_name: str + ) -> int: + """Count intersecting values between two columns.""" + data1 = self.check_data(table1.data) + data2 = self.check_data(table2.data) + fqn1 = self._get_fqn(data1.identifier) + fqn2 = self._get_fqn(data2.identifier) + + query = f""" + SELECT COUNT(*) as count + FROM ( + SELECT DISTINCT `{column1_name}` as key + FROM {fqn1} + WHERE `{column1_name}` IS NOT NULL + ) t1 + INNER JOIN ( + SELECT DISTINCT `{column2_name}` as key + FROM {fqn2} + WHERE `{column2_name}` IS NOT NULL + ) t2 + ON t1.key = t2.key + """ + return self._execute_sql(query)[0]["count"] + + def get_composite_key_uniqueness( + self, table_name: str, columns: list[str], dataset_data: DataSetData + ) -> int: + """Get the number of unique composite key combinations.""" + data = self.check_data(dataset_data) + fqn = self._get_fqn(data.identifier) + + # Build column list with backticks + safe_columns = [f"`{col}`" for col in columns] + columns_str = ", ".join(safe_columns) + + # Build null filter + null_filter = " AND ".join(f"{col} IS NOT NULL" for col in safe_columns) + + query = f""" + SELECT COUNT(*) as count + FROM ( + SELECT DISTINCT {columns_str} + FROM {fqn} + WHERE {null_filter} + ) + """ + return self._execute_sql(query)[0]["count"] + + def intersect_composite_keys_count( + self, + table1: "DataSet", + columns1: list[str], + table2: "DataSet", + columns2: list[str], + ) -> int: + """Count intersecting composite key combinations between two tables.""" + if len(columns1) != len(columns2): + raise ValueError("Column lists must have the same length for composite key intersection.") + + data1 = self.check_data(table1.data) + data2 = self.check_data(table2.data) + fqn1 = self._get_fqn(data1.identifier) + fqn2 = self._get_fqn(data2.identifier) + + # Build column lists with backticks + safe_columns1 = [f"`{col}`" for col in columns1] + safe_columns2 = [f"`{col}`" for col in columns2] + + # Subquery for distinct keys from table 1 + distinct_cols1 = ", ".join(safe_columns1) + null_filter1 = " AND ".join(f"{c} IS NOT NULL" for c in safe_columns1) + subquery1 = f"""( + SELECT DISTINCT {distinct_cols1} + FROM {fqn1} + WHERE {null_filter1} + ) AS t1""" + + # Subquery for distinct keys from table 2 + distinct_cols2 = ", ".join(safe_columns2) + null_filter2 = " AND ".join(f"{c} IS NOT NULL" for c in safe_columns2) + subquery2 = f"""( + SELECT DISTINCT {distinct_cols2} + FROM {fqn2} + WHERE {null_filter2} + ) AS t2""" + + # Join conditions + join_conditions = " AND ".join( + [f"t1.{c1} = t2.{c2}" for c1, c2 in zip(safe_columns1, safe_columns2)] + ) + + query = f""" + SELECT COUNT(*) as count + FROM {subquery1} + INNER JOIN {subquery2} ON {join_conditions} + """ + return self._execute_sql(query)[0]["count"] + + def get_details(self, data: BigQueryConfig): + """Get configuration details.""" + data = self.check_data(data) + return data.model_dump() + + +def can_handle_bigquery(df: Any) -> bool: + """Check if the data is a BigQuery config.""" + try: + BigQueryConfig.model_validate(df) + return True + except Exception: + return False + + +def register(factory: AdapterFactory): + """Register the BigQuery adapter with the factory.""" + if BIGQUERY_AVAILABLE: + factory.register("bigquery", can_handle_bigquery, BigQueryAdapter, BigQueryConfig) diff --git a/src/intugle/adapters/types/bigquery/models.py b/src/intugle/adapters/types/bigquery/models.py new file mode 100644 index 0000000..87f4f3c --- /dev/null +++ b/src/intugle/adapters/types/bigquery/models.py @@ -0,0 +1,17 @@ +from typing import Literal, Optional + +from pydantic import Field + +from intugle.common.schema import SchemaBase + + +class BigQueryConnectionConfig(SchemaBase): + project_id: str + dataset_id: str = Field(..., alias="dataset") + credentials_path: Optional[str] = None + location: str = "US" + + +class BigQueryConfig(SchemaBase): + identifier: str + type: Literal["bigquery"] = "bigquery" diff --git a/tests/adapters/test_bigquery_adapter.py b/tests/adapters/test_bigquery_adapter.py new file mode 100644 index 0000000..b416620 --- /dev/null +++ b/tests/adapters/test_bigquery_adapter.py @@ -0,0 +1,262 @@ +import pytest +from unittest.mock import MagicMock, patch, PropertyMock +from intugle.adapters.types.bigquery.bigquery import ( + BigQueryAdapter, + can_handle_bigquery, + BIGQUERY_AVAILABLE, +) +from intugle.adapters.types.bigquery.models import BigQueryConfig, BigQueryConnectionConfig + + +@pytest.fixture +def mock_bigquery_client(): + """Mock BigQuery client.""" + with patch("intugle.adapters.types.bigquery.bigquery.bigquery") as mock_bq: + mock_client = MagicMock() + mock_bq.Client.return_value = mock_client + yield mock_client + + +@pytest.fixture +def mock_settings(): + """Mock settings with BigQuery configuration.""" + with patch("intugle.adapters.types.bigquery.bigquery.settings") as mock_settings: + mock_settings.PROFILES = { + "bigquery": { + "project_id": "test-project", + "dataset": "test_dataset", + "name": "test_bigquery_source", + "location": "US", + } + } + yield mock_settings + + +@pytest.fixture +def bigquery_config(): + """Create a test BigQuery config.""" + return BigQueryConfig(identifier="test_table", type="bigquery") + + +@pytest.mark.skipif(not BIGQUERY_AVAILABLE, reason="BigQuery dependencies not installed") +class TestBigQueryAdapterContract: + """Test that BigQueryAdapter implements the Adapter interface correctly.""" + + def test_implements_required_methods(self, mock_settings, mock_bigquery_client): + """Test that all abstract methods are implemented.""" + adapter = BigQueryAdapter() + + # Check all abstract methods exist + assert hasattr(adapter, "profile") + assert hasattr(adapter, "column_profile") + assert hasattr(adapter, "load") + assert hasattr(adapter, "execute") + assert hasattr(adapter, "to_df") + assert hasattr(adapter, "to_df_from_query") + assert hasattr(adapter, "create_table_from_query") + assert hasattr(adapter, "create_new_config_from_etl") + assert hasattr(adapter, "intersect_count") + assert hasattr(adapter, "get_composite_key_uniqueness") + assert hasattr(adapter, "intersect_composite_keys_count") + assert callable(adapter.profile) + assert callable(adapter.column_profile) + + def test_inherits_from_adapter(self, mock_settings, mock_bigquery_client): + """Test that BigQueryAdapter inherits from Adapter base class.""" + from intugle.adapters.adapter import Adapter + + adapter = BigQueryAdapter() + assert isinstance(adapter, Adapter) + + def test_singleton_pattern(self, mock_settings, mock_bigquery_client): + """Test that BigQueryAdapter follows singleton pattern.""" + # Reset singleton + BigQueryAdapter._instance = None + BigQueryAdapter._initialized = False + + adapter1 = BigQueryAdapter() + adapter2 = BigQueryAdapter() + assert adapter1 is adapter2 + + def test_registered_in_factory(self): + """Test that BigQueryAdapter is registered in the adapter factory.""" + from intugle.adapters.factory import AdapterFactory + + factory = AdapterFactory() + assert "bigquery" in factory.dataframe_funcs + + +@pytest.mark.skipif(not BIGQUERY_AVAILABLE, reason="BigQuery dependencies not installed") +class TestBigQuerySpecificBehavior: + """Test BigQuery-specific functionality.""" + + def test_database_and_schema_properties(self, mock_settings, mock_bigquery_client): + """Test that database and schema properties work correctly.""" + BigQueryAdapter._instance = None + BigQueryAdapter._initialized = False + + adapter = BigQueryAdapter() + assert adapter.database == "test-project" + assert adapter.schema == "test_dataset" + + adapter.database = "new-project" + adapter.schema = "new_dataset" + assert adapter.database == "new-project" + assert adapter.schema == "new_dataset" + + def test_check_data_validates_bigquery_config(self, mock_settings, mock_bigquery_client, bigquery_config): + """Test that check_data validates BigQuery config correctly.""" + BigQueryAdapter._instance = None + BigQueryAdapter._initialized = False + + adapter = BigQueryAdapter() + validated = adapter.check_data(bigquery_config) + assert isinstance(validated, BigQueryConfig) + assert validated.identifier == "test_table" + assert validated.type == "bigquery" + + def test_check_data_rejects_invalid_config(self, mock_settings, mock_bigquery_client): + """Test that check_data rejects invalid configs.""" + BigQueryAdapter._instance = None + BigQueryAdapter._initialized = False + + adapter = BigQueryAdapter() + with pytest.raises(TypeError, match="Input must be a BigQuery config"): + adapter.check_data({"invalid": "config"}) + + def test_create_new_config_from_etl(self, mock_settings, mock_bigquery_client): + """Test creating a new config from ETL name.""" + BigQueryAdapter._instance = None + BigQueryAdapter._initialized = False + + adapter = BigQueryAdapter() + config = adapter.create_new_config_from_etl("etl_table") + assert isinstance(config, BigQueryConfig) + assert config.identifier == "etl_table" + assert config.type == "bigquery" + + def test_get_details_returns_config_dict(self, mock_settings, mock_bigquery_client, bigquery_config): + """Test that get_details returns config as dictionary.""" + BigQueryAdapter._instance = None + BigQueryAdapter._initialized = False + + adapter = BigQueryAdapter() + details = adapter.get_details(bigquery_config) + assert isinstance(details, dict) + assert details["identifier"] == "test_table" + assert details["type"] == "bigquery" + + def test_get_fqn_creates_fully_qualified_name(self, mock_settings, mock_bigquery_client): + """Test that _get_fqn creates proper fully qualified names.""" + BigQueryAdapter._instance = None + BigQueryAdapter._initialized = False + + adapter = BigQueryAdapter() + + # Simple table name + fqn = adapter._get_fqn("test_table") + assert fqn == "`test-project.test_dataset.test_table`" + + # Dataset.table format + fqn = adapter._get_fqn("other_dataset.test_table") + assert fqn == "`test-project.other_dataset.test_table`" + + # Full project.dataset.table format + fqn = adapter._get_fqn("other-project.other_dataset.test_table") + assert fqn == "`other-project.other_dataset.test_table`" + + def test_client_property_exists(self, mock_settings, mock_bigquery_client): + """Test that the BigQuery client is initialized.""" + BigQueryAdapter._instance = None + BigQueryAdapter._initialized = False + + adapter = BigQueryAdapter() + assert adapter.client is not None + + +@pytest.mark.skipif(not BIGQUERY_AVAILABLE, reason="BigQuery dependencies not installed") +class TestCanHandleBigQuery: + """Test the can_handle_bigquery function.""" + + def test_handles_bigquery_config_object(self, bigquery_config): + """Test that it accepts valid BigQuery config objects.""" + assert can_handle_bigquery(bigquery_config) is True + + def test_handles_valid_bigquery_dict(self): + """Test that it accepts valid BigQuery config dictionaries.""" + config_dict = {"identifier": "test_table", "type": "bigquery"} + assert can_handle_bigquery(config_dict) is True + + def test_rejects_non_bigquery_configs(self): + """Test that it rejects non-BigQuery configs.""" + assert can_handle_bigquery({"type": "postgres"}) is False + assert can_handle_bigquery({"invalid": "config"}) is False + assert can_handle_bigquery("not a config") is False + assert can_handle_bigquery(None) is False + + +@pytest.mark.skipif(not BIGQUERY_AVAILABLE, reason="BigQuery dependencies not installed") +class TestBigQueryErrorHandling: + """Test error handling in BigQueryAdapter.""" + + def test_import_error_when_bigquery_not_installed(self): + """Test that ImportError is raised when BigQuery is not installed.""" + with patch("intugle.adapters.types.bigquery.bigquery.BIGQUERY_AVAILABLE", False): + BigQueryAdapter._instance = None + BigQueryAdapter._initialized = False + + with pytest.raises(ImportError, match="BigQuery dependencies are not installed"): + BigQueryAdapter() + + def test_check_data_type_error_message(self, mock_settings, mock_bigquery_client): + """Test that check_data provides clear error message.""" + BigQueryAdapter._instance = None + BigQueryAdapter._initialized = False + + adapter = BigQueryAdapter() + with pytest.raises(TypeError) as exc_info: + adapter.check_data({"not": "valid"}) + assert "Input must be a BigQuery config" in str(exc_info.value) + + def test_connection_error_provides_context(self, mock_bigquery_client): + """Test that connection errors provide helpful context.""" + with patch("intugle.adapters.types.bigquery.bigquery.settings") as mock_settings: + mock_settings.PROFILES = {} + BigQueryAdapter._instance = None + BigQueryAdapter._initialized = False + + with pytest.raises(ValueError, match="No 'bigquery' section found in profiles.yml"): + BigQueryAdapter() + + +@pytest.mark.skipif(not BIGQUERY_AVAILABLE, reason="BigQuery dependencies not installed") +class TestBigQueryConnectionConfig: + """Test BigQueryConnectionConfig model.""" + + def test_valid_config(self): + """Test that valid config can be created.""" + config = BigQueryConnectionConfig( + project_id="test-project", + dataset="test_dataset", + location="US", + ) + assert config.project_id == "test-project" + assert config.dataset_id == "test_dataset" + assert config.location == "US" + + def test_optional_credentials_path(self): + """Test that credentials_path is optional.""" + config = BigQueryConnectionConfig( + project_id="test-project", + dataset="test_dataset", + ) + assert config.credentials_path is None + + def test_with_credentials_path(self): + """Test config with credentials path.""" + config = BigQueryConnectionConfig( + project_id="test-project", + dataset="test_dataset", + credentials_path="/path/to/credentials.json", + ) + assert config.credentials_path == "/path/to/credentials.json"