mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-25 01:43:27 +08:00
Feature big query connector (#15871)
### What problem does this PR solve? This PR adds Google BigQuery as a first-class data source connector in RAGFlow. It enables users to ingest and sync BigQuery data using the same row-to-document model used by relational database connectors: selected content columns become document text, metadata columns become document metadata, an optional ID column provides stable document IDs, and an optional timestamp column enables cursor-based incremental sync. The connector supports service-account JSON credentials, table mode, custom query mode, GoogleSQL queries, cursor-based incremental sync, deleted-row pruning support, configurable query limits such as `maximum_bytes_billed`, dry-run validation, batch loading, stable document IDs, and BigQuery-aware value serialization.
This commit is contained in:
Attili-sys
430
test/unit_test/data_source/test_bigquery_connector.py
Normal file
430
test/unit_test/data_source/test_bigquery_connector.py
Normal file
@@ -0,0 +1,430 @@
|
||||
#
|
||||
# Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
import json
|
||||
import types
|
||||
from datetime import date, datetime, time, timezone
|
||||
from decimal import Decimal
|
||||
|
||||
import pytest
|
||||
|
||||
from common.data_source import bigquery_connector as bq_mod
|
||||
from common.data_source.bigquery_connector import BigQueryConnector
|
||||
from common.data_source.config import DocumentSource
|
||||
from common.data_source.exceptions import (
|
||||
ConnectorMissingCredentialError,
|
||||
ConnectorValidationError,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fakes for the google-cloud-bigquery client surface
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeSchemaField:
|
||||
def __init__(self, name, field_type):
|
||||
self.name = name
|
||||
self.field_type = field_type
|
||||
|
||||
|
||||
class _FakeScalarQueryParameter:
|
||||
def __init__(self, name, type_, value):
|
||||
self.name = name
|
||||
self.type_ = type_
|
||||
self.value = value
|
||||
|
||||
|
||||
class _FakeQueryJobConfig:
|
||||
def __init__(self):
|
||||
self.use_legacy_sql = None
|
||||
self.use_query_cache = None
|
||||
self.maximum_bytes_billed = None
|
||||
self.job_timeout_ms = None
|
||||
self.query_parameters = None
|
||||
self.dry_run = False
|
||||
|
||||
|
||||
class _FakeJob:
|
||||
def __init__(self, rows=None, schema=None, total_bytes=123):
|
||||
self._rows = rows if rows is not None else []
|
||||
self.schema = schema or []
|
||||
self.total_bytes_processed = total_bytes
|
||||
|
||||
def result(self, page_size=None):
|
||||
return iter(self._rows)
|
||||
|
||||
|
||||
class _FakeClient:
|
||||
def __init__(self, rows=None, result_schema=None, table_schema=None):
|
||||
self._rows = rows if rows is not None else []
|
||||
self._result_schema = result_schema or []
|
||||
self._table_schema = table_schema or []
|
||||
self.queries = []
|
||||
|
||||
def query(self, query, job_config=None, location=None):
|
||||
self.queries.append((query, job_config, location))
|
||||
return _FakeJob(rows=self._rows, schema=self._result_schema)
|
||||
|
||||
def get_table(self, ref):
|
||||
return types.SimpleNamespace(schema=self._table_schema)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _patch_bigquery(monkeypatch):
|
||||
fake = types.SimpleNamespace(
|
||||
QueryJobConfig=_FakeQueryJobConfig,
|
||||
ScalarQueryParameter=_FakeScalarQueryParameter,
|
||||
)
|
||||
monkeypatch.setattr(bq_mod, "bigquery", fake)
|
||||
monkeypatch.setattr(bq_mod, "service_account", types.SimpleNamespace())
|
||||
|
||||
|
||||
SERVICE_ACCOUNT = {"type": "service_account", "project_id": "p"}
|
||||
|
||||
|
||||
def _make_connector(client=None, **overrides):
|
||||
kwargs = dict(
|
||||
project_id="my-proj",
|
||||
dataset_id="ds",
|
||||
table_id="tbl",
|
||||
content_columns="name,description",
|
||||
metadata_columns="status",
|
||||
id_column="id",
|
||||
timestamp_column=None,
|
||||
batch_size=2,
|
||||
)
|
||||
kwargs.update(overrides)
|
||||
connector = BigQueryConnector(**kwargs)
|
||||
connector._credentials = {"service_account_info": SERVICE_ACCOUNT}
|
||||
if client is not None:
|
||||
connector._client = client
|
||||
return connector
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Credentials
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_load_credentials_accepts_json_string():
|
||||
connector = BigQueryConnector(project_id="p", content_columns="name")
|
||||
connector.load_credentials({"service_account_json": json.dumps(SERVICE_ACCOUNT)})
|
||||
assert connector._credentials["service_account_info"] == SERVICE_ACCOUNT
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_load_credentials_accepts_dict():
|
||||
connector = BigQueryConnector(project_id="p", content_columns="name")
|
||||
connector.load_credentials({"service_account_json": SERVICE_ACCOUNT})
|
||||
assert connector._credentials["service_account_info"] == SERVICE_ACCOUNT
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_missing_credentials_raises():
|
||||
connector = BigQueryConnector(project_id="p", content_columns="name")
|
||||
with pytest.raises(ConnectorMissingCredentialError):
|
||||
connector.load_credentials({})
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_invalid_credentials_json_raises():
|
||||
connector = BigQueryConnector(project_id="p", content_columns="name")
|
||||
with pytest.raises(ConnectorMissingCredentialError):
|
||||
connector.load_credentials({"service_account_json": "{not-json"})
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Query construction
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_table_query_quotes_identifiers():
|
||||
connector = _make_connector()
|
||||
assert connector._build_base_query() == "SELECT * FROM `my-proj.ds.tbl`"
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_custom_query_strips_trailing_semicolon():
|
||||
connector = _make_connector(query="SELECT * FROM t WHERE x = 1;")
|
||||
assert connector._build_base_query() == "SELECT * FROM t WHERE x = 1"
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_missing_table_and_query_raises():
|
||||
connector = _make_connector(dataset_id="", table_id="", query="")
|
||||
with pytest.raises(ConnectorValidationError):
|
||||
connector._build_base_query()
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_time_filtered_query_compound_cursor_with_id_column():
|
||||
client = _FakeClient(table_schema=[
|
||||
_FakeSchemaField("updated_at", "TIMESTAMP"),
|
||||
_FakeSchemaField("id", "STRING")
|
||||
])
|
||||
connector = _make_connector(client=client, timestamp_column="updated_at", id_column="id")
|
||||
start = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
end = datetime(2026, 2, 1, tzinfo=timezone.utc)
|
||||
|
||||
query, params = connector._build_time_filtered_query(
|
||||
connector._build_base_query(), start, end, start_id="last-id"
|
||||
)
|
||||
|
||||
assert "(ragflow_src.updated_at > @start_cursor OR (ragflow_src.updated_at = @start_cursor AND ragflow_src.id > @start_cursor_id))" in query
|
||||
assert "ragflow_src.updated_at <= @end_cursor" in query
|
||||
assert [(p.name, p.type_, p.value) for p in params] == [
|
||||
("start_cursor", "TIMESTAMP", start),
|
||||
("start_cursor_id", "STRING", "last-id"),
|
||||
("end_cursor", "TIMESTAMP", end),
|
||||
]
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_time_filtered_query_uses_gte_without_id_column():
|
||||
client = _FakeClient(table_schema=[_FakeSchemaField("updated_at", "TIMESTAMP")])
|
||||
connector = _make_connector(client=client, timestamp_column="updated_at", id_column=None)
|
||||
start = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
end = datetime(2026, 2, 1, tzinfo=timezone.utc)
|
||||
|
||||
query, params = connector._build_time_filtered_query(
|
||||
connector._build_base_query(), start, end
|
||||
)
|
||||
|
||||
assert "ragflow_src.updated_at >= @start_cursor" in query
|
||||
assert "ragflow_src.updated_at <= @end_cursor" in query
|
||||
assert [(p.name, p.type_, p.value) for p in params] == [
|
||||
("start_cursor", "TIMESTAMP", start),
|
||||
("end_cursor", "TIMESTAMP", end),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_cursor_param_type_resolves_for_date_and_int():
|
||||
date_client = _FakeClient(table_schema=[_FakeSchemaField("d", "DATE")])
|
||||
int_client = _FakeClient(table_schema=[_FakeSchemaField("n", "INTEGER")])
|
||||
|
||||
date_conn = _make_connector(client=date_client, timestamp_column="d")
|
||||
int_conn = _make_connector(client=int_client, timestamp_column="n")
|
||||
|
||||
assert date_conn._resolve_cursor_param_type() == "DATE"
|
||||
assert int_conn._resolve_cursor_param_type() == "INT64"
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_unsupported_cursor_type_raises():
|
||||
client = _FakeClient(table_schema=[_FakeSchemaField("ts", "BOOL")])
|
||||
connector = _make_connector(client=client, timestamp_column="ts")
|
||||
with pytest.raises(ConnectorValidationError):
|
||||
connector._resolve_cursor_param_type()
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_slim_query_uses_id_column():
|
||||
connector = _make_connector()
|
||||
slim = connector._build_slim_query(connector._build_base_query())
|
||||
assert "ragflow_src.id" in slim
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Value rendering & row conversion
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_value_serialization_types():
|
||||
assert BigQueryConnector._render_content_value(Decimal("1.50")) == "1.50"
|
||||
assert BigQueryConnector._render_content_value(b"binary") is None
|
||||
assert BigQueryConnector._render_content_value(date(2026, 1, 2)) == "2026-01-02"
|
||||
assert BigQueryConnector._render_content_value(time(3, 4, 5)) == "03:04:05"
|
||||
assert (
|
||||
BigQueryConnector._render_content_value({"a": 1, "b": [1, 2]})
|
||||
== '{"a": 1, "b": [1, 2]}'
|
||||
)
|
||||
assert (
|
||||
BigQueryConnector._render_content_value("POINT(1 2)") == "POINT(1 2)"
|
||||
) # GEOGRAPHY WKT passes through
|
||||
|
||||
# Metadata base64-encodes bytes instead of skipping.
|
||||
assert BigQueryConnector._render_metadata_value(b"hi") == "aGk="
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_row_to_document_content_metadata_id_timestamp():
|
||||
connector = _make_connector(timestamp_column="updated_at")
|
||||
ts = datetime(2026, 3, 4, 5, 6, 7, tzinfo=timezone.utc)
|
||||
row = {
|
||||
"id": 42,
|
||||
"name": "Acme",
|
||||
"description": "A company",
|
||||
"status": "active",
|
||||
"updated_at": ts,
|
||||
}
|
||||
|
||||
doc = connector._row_to_document(row)
|
||||
|
||||
assert doc.id == "bigquery:my-proj:ds.tbl:42"
|
||||
assert "【name】:\nAcme" in doc.blob.decode("utf-8")
|
||||
assert "【description】:\nA company" in doc.blob.decode("utf-8")
|
||||
assert doc.metadata == {"status": "active"}
|
||||
assert doc.source == DocumentSource.BIGQUERY
|
||||
assert doc.extension == ".txt"
|
||||
assert doc.doc_updated_at == ts
|
||||
assert doc.semantic_identifier == "Acme"
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_document_id_falls_back_to_content_hash_without_id_column():
|
||||
connector = _make_connector(id_column=None)
|
||||
row = {"name": "Acme", "description": "A company", "status": "active"}
|
||||
doc_id = connector._build_document_id_from_row(row)
|
||||
assert doc_id.startswith("bigquery:my-proj:ds.tbl:")
|
||||
assert len(doc_id.rsplit(":", 1)[1]) == 32 # md5 hex
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_custom_query_mode_id_prefix():
|
||||
connector = _make_connector(dataset_id="", table_id="", query="SELECT * FROM t")
|
||||
row = {"id": 7, "name": "x"}
|
||||
assert connector._build_document_id_from_row(row) == "bigquery:my-proj:query:7"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Batching & cursor round-trip
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_batches_accumulate_to_batch_size():
|
||||
rows = [
|
||||
{"id": i, "name": f"n{i}", "description": "d", "status": "s"}
|
||||
for i in range(5)
|
||||
]
|
||||
client = _FakeClient(rows=rows)
|
||||
connector = _make_connector(client=client, batch_size=2)
|
||||
|
||||
batches = list(connector.load_from_state())
|
||||
|
||||
assert [len(b) for b in batches] == [2, 2, 1]
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_cursor_serialize_deserialize_roundtrip():
|
||||
dt = datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)
|
||||
d = date(2026, 1, 1)
|
||||
t = time(12, 0)
|
||||
dec = Decimal("1.23")
|
||||
|
||||
assert BigQueryConnector.deserialize_cursor_value(
|
||||
BigQueryConnector.serialize_cursor_value(dt)
|
||||
) == dt
|
||||
assert BigQueryConnector.deserialize_cursor_value(
|
||||
BigQueryConnector.serialize_cursor_value(d)
|
||||
) == d
|
||||
assert BigQueryConnector.deserialize_cursor_value(
|
||||
BigQueryConnector.serialize_cursor_value(t)
|
||||
) == t
|
||||
assert BigQueryConnector.deserialize_cursor_value(
|
||||
BigQueryConnector.serialize_cursor_value(dec)
|
||||
) == dec
|
||||
assert BigQueryConnector.serialize_cursor_value(42) == 42
|
||||
assert BigQueryConnector.deserialize_cursor_value(42) == 42
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_load_from_cursor_range_empty_without_end():
|
||||
connector = _make_connector(timestamp_column="updated_at")
|
||||
assert list(connector.load_from_cursor_range(start_value=None, end_value=None)) == []
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_load_from_cursor_range_empty_when_end_not_after_start():
|
||||
connector = _make_connector(timestamp_column="updated_at")
|
||||
start = datetime(2026, 2, 1, tzinfo=timezone.utc)
|
||||
end = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
assert list(connector.load_from_cursor_range(start, end)) == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Slim docs & validation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_slim_docs_use_id_column():
|
||||
rows = [{"id": 1}, {"id": 2}]
|
||||
client = _FakeClient(rows=rows)
|
||||
connector = _make_connector(client=client, batch_size=10)
|
||||
|
||||
slim_batches = list(connector.retrieve_all_slim_docs_perm_sync())
|
||||
ids = [doc.id for batch in slim_batches for doc in batch]
|
||||
|
||||
assert ids == ["bigquery:my-proj:ds.tbl:1", "bigquery:my-proj:ds.tbl:2"]
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_validation_detects_missing_content_column():
|
||||
client = _FakeClient(table_schema=[_FakeSchemaField("other", "STRING")])
|
||||
connector = _make_connector(client=client, content_columns="name")
|
||||
with pytest.raises(ConnectorValidationError, match="name"):
|
||||
connector.validate_connector_settings()
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_validation_detects_missing_metadata_column():
|
||||
client = _FakeClient(table_schema=[_FakeSchemaField("name", "STRING")])
|
||||
connector = _make_connector(client=client, content_columns="name", metadata_columns="status")
|
||||
with pytest.raises(ConnectorValidationError, match="status"):
|
||||
connector.validate_connector_settings()
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_validation_detects_missing_id_column():
|
||||
client = _FakeClient(table_schema=[_FakeSchemaField("name", "STRING")])
|
||||
connector = _make_connector(client=client, content_columns="name", metadata_columns="", id_column="id")
|
||||
with pytest.raises(ConnectorValidationError, match="id"):
|
||||
connector.validate_connector_settings()
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_validation_detects_missing_timestamp_column():
|
||||
client = _FakeClient(table_schema=[_FakeSchemaField("name", "STRING")])
|
||||
connector = _make_connector(client=client, content_columns="name", metadata_columns="", id_column="", timestamp_column="ts")
|
||||
with pytest.raises(ConnectorValidationError, match="ts"):
|
||||
connector.validate_connector_settings()
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_validation_detects_unsupported_cursor_type_early():
|
||||
client = _FakeClient(table_schema=[_FakeSchemaField("name", "STRING"), _FakeSchemaField("ts", "BOOL")])
|
||||
connector = _make_connector(client=client, content_columns="name", metadata_columns="", id_column="", timestamp_column="ts")
|
||||
with pytest.raises(ConnectorValidationError, match="not supported as a cursor"):
|
||||
connector.validate_connector_settings()
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_validation_missing_credentials_raises():
|
||||
connector = BigQueryConnector(project_id="p", content_columns="name")
|
||||
with pytest.raises(ConnectorMissingCredentialError):
|
||||
connector.validate_connector_settings()
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_validation_dry_run_failure_becomes_validation_error():
|
||||
class _FailingClient(_FakeClient):
|
||||
def query(self, query, job_config=None, location=None):
|
||||
raise RuntimeError("bad query / would scan too much")
|
||||
|
||||
connector = _make_connector(client=_FailingClient())
|
||||
with pytest.raises(ConnectorValidationError):
|
||||
connector.validate_connector_settings()
|
||||
@@ -314,7 +314,7 @@ class _FakeRDBMSConnector:
|
||||
self.load_from_state_called = True
|
||||
return iter((["full-sync"],))
|
||||
|
||||
def load_from_cursor_range(self, start_value=None, end_value=None):
|
||||
def load_from_cursor_range(self, start_value=None, start_id=None, end_value=None):
|
||||
self.load_from_cursor_range_called = True
|
||||
return iter(([ _make_fake_doc("incremental-doc") ],))
|
||||
|
||||
@@ -405,6 +405,54 @@ async def test_rdbms_cursor_persists_only_after_success(monkeypatch):
|
||||
assert connector.persist_sync_state_called is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.p2
|
||||
async def test_rdbms_cursor_does_not_persist_when_parse_returns_errors(monkeypatch):
|
||||
monkeypatch.setattr(sync_data_source, "RDBMSConnector", _FakeRDBMSConnector)
|
||||
_patch_common_dependencies(monkeypatch)
|
||||
monkeypatch.setattr(
|
||||
sync_data_source.KnowledgebaseService,
|
||||
"get_by_id",
|
||||
lambda *_args, **_kwargs: (True, object()),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
sync_data_source.SyncLogsService,
|
||||
"increase_docs",
|
||||
lambda *_args, **_kwargs: None,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
sync_data_source.SyncLogsService,
|
||||
"duplicate_and_parse",
|
||||
lambda *_args, **_kwargs: (["parse error"], ["parsed-doc-id"]),
|
||||
)
|
||||
|
||||
task = {
|
||||
**_make_task(),
|
||||
"reindex": "0",
|
||||
"poll_range_start": datetime(2026, 1, 1, tzinfo=timezone.utc),
|
||||
"skip_connection_log": True,
|
||||
}
|
||||
sync = sync_data_source.MySQL(
|
||||
{
|
||||
"host": "localhost",
|
||||
"port": 3306,
|
||||
"database": "db",
|
||||
"query": "SELECT * FROM t",
|
||||
"content_columns": "name",
|
||||
"timestamp_column": "ts",
|
||||
"credentials": {"username": "u", "password": "p"},
|
||||
"sync_deleted_files": False,
|
||||
}
|
||||
)
|
||||
|
||||
await sync._run_task_logic(task)
|
||||
|
||||
connector = _FakeRDBMSConnector.instance
|
||||
assert connector is not None
|
||||
assert connector.persist_sync_state_called is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.p2
|
||||
async def test_rdbms_cursor_does_not_persist_when_batch_is_skipped(monkeypatch):
|
||||
@@ -456,6 +504,263 @@ async def test_rdbms_cursor_does_not_persist_when_batch_is_skipped(monkeypatch):
|
||||
assert connector.persist_sync_state_called is False
|
||||
|
||||
|
||||
class _FakeBigQueryConnector:
|
||||
instance = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
project_id,
|
||||
dataset_id=None,
|
||||
table_id=None,
|
||||
location=None,
|
||||
query="",
|
||||
content_columns="",
|
||||
metadata_columns=None,
|
||||
id_column=None,
|
||||
timestamp_column=None,
|
||||
batch_size=2,
|
||||
page_size=1000,
|
||||
maximum_bytes_billed=None,
|
||||
job_timeout_ms=None,
|
||||
use_query_cache=True,
|
||||
):
|
||||
self.project_id = project_id
|
||||
self.dataset_id = dataset_id
|
||||
self.table_id = table_id
|
||||
self.query = query
|
||||
self.content_columns = content_columns
|
||||
self.timestamp_column = timestamp_column
|
||||
self.batch_size = batch_size
|
||||
self.load_from_state_called = False
|
||||
self.load_from_cursor_range_called = False
|
||||
self.retrieve_all_slim_docs_perm_sync_called = False
|
||||
self.prepare_sync_state_called = False
|
||||
self.persist_sync_state_called = False
|
||||
self._pending_sync_cursor_value = None
|
||||
_FakeBigQueryConnector.instance = self
|
||||
|
||||
def load_credentials(self, credentials):
|
||||
self.credentials = credentials
|
||||
|
||||
def validate_connector_settings(self):
|
||||
return None
|
||||
|
||||
def prepare_sync_state(self, connector_id, config):
|
||||
self.prepare_sync_state_called = True
|
||||
self.prepare_sync_state_args = (connector_id, config)
|
||||
|
||||
def get_saved_sync_cursor_value(self):
|
||||
return None
|
||||
|
||||
def retrieve_all_slim_docs_perm_sync(self, callback=None):
|
||||
del callback
|
||||
self.retrieve_all_slim_docs_perm_sync_called = True
|
||||
yield [types.SimpleNamespace(id="bq-row-1")]
|
||||
|
||||
def load_from_state(self):
|
||||
self.load_from_state_called = True
|
||||
return iter((["full-sync"],))
|
||||
|
||||
def load_from_cursor_range(self, start_value=None, start_id=None, end_value=None):
|
||||
self.load_from_cursor_range_called = True
|
||||
return iter(([_make_fake_doc("bq-incremental-doc")],))
|
||||
|
||||
def persist_sync_state(self):
|
||||
self.persist_sync_state_called = True
|
||||
|
||||
|
||||
def _bigquery_conf(**overrides):
|
||||
conf = {
|
||||
"project_id": "proj",
|
||||
"dataset_id": "ds",
|
||||
"table_id": "tbl",
|
||||
"content_columns": "name",
|
||||
"credentials": {"service_account_json": "{}"},
|
||||
"sync_deleted_files": False,
|
||||
}
|
||||
conf.update(overrides)
|
||||
return conf
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.p2
|
||||
async def test_bigquery_generate_full_sync_on_first_run(monkeypatch):
|
||||
monkeypatch.setattr(sync_data_source, "BigQueryConnector", _FakeBigQueryConnector)
|
||||
|
||||
task = {
|
||||
**_make_task(),
|
||||
"reindex": "0",
|
||||
"poll_range_start": None,
|
||||
"skip_connection_log": True,
|
||||
}
|
||||
sync = sync_data_source.BigQuery(_bigquery_conf())
|
||||
|
||||
document_generator = await sync._generate(task)
|
||||
connector = _FakeBigQueryConnector.instance
|
||||
|
||||
assert connector is not None
|
||||
assert connector.prepare_sync_state_called is True
|
||||
assert connector.load_from_state_called is True
|
||||
assert connector.load_from_cursor_range_called is False
|
||||
assert list(document_generator) == [["full-sync"]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.p2
|
||||
async def test_bigquery_generate_incremental_cursor_path(monkeypatch):
|
||||
monkeypatch.setattr(sync_data_source, "BigQueryConnector", _FakeBigQueryConnector)
|
||||
|
||||
task = {
|
||||
**_make_task(),
|
||||
"reindex": "0",
|
||||
"poll_range_start": datetime(2026, 1, 1, tzinfo=timezone.utc),
|
||||
"skip_connection_log": True,
|
||||
}
|
||||
sync = sync_data_source.BigQuery(_bigquery_conf(timestamp_column="updated_at"))
|
||||
|
||||
document_generator = await sync._generate(task)
|
||||
connector = _FakeBigQueryConnector.instance
|
||||
|
||||
assert connector is not None
|
||||
assert connector.load_from_cursor_range_called is True
|
||||
assert connector.load_from_state_called is False
|
||||
assert [doc.id for doc in list(document_generator)[0]] == ["bq-incremental-doc"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.p2
|
||||
async def test_bigquery_cursor_persists_only_after_success(monkeypatch):
|
||||
monkeypatch.setattr(sync_data_source, "BigQueryConnector", _FakeBigQueryConnector)
|
||||
_patch_common_dependencies(monkeypatch)
|
||||
monkeypatch.setattr(
|
||||
sync_data_source.KnowledgebaseService,
|
||||
"get_by_id",
|
||||
lambda *_args, **_kwargs: (True, object()),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
sync_data_source.SyncLogsService,
|
||||
"increase_docs",
|
||||
lambda *_args, **_kwargs: None,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
sync_data_source.SyncLogsService,
|
||||
"duplicate_and_parse",
|
||||
lambda *_args, **_kwargs: ([], ["parsed-doc-id"]),
|
||||
)
|
||||
|
||||
task = {
|
||||
**_make_task(),
|
||||
"reindex": "0",
|
||||
"poll_range_start": datetime(2026, 1, 1, tzinfo=timezone.utc),
|
||||
"skip_connection_log": True,
|
||||
}
|
||||
sync = sync_data_source.BigQuery(_bigquery_conf(timestamp_column="updated_at"))
|
||||
|
||||
await sync._run_task_logic(task)
|
||||
|
||||
connector = _FakeBigQueryConnector.instance
|
||||
assert connector is not None
|
||||
assert connector.persist_sync_state_called is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.p2
|
||||
async def test_bigquery_cursor_does_not_persist_when_parse_returns_errors(monkeypatch):
|
||||
monkeypatch.setattr(sync_data_source, "BigQueryConnector", _FakeBigQueryConnector)
|
||||
_patch_common_dependencies(monkeypatch)
|
||||
monkeypatch.setattr(
|
||||
sync_data_source.KnowledgebaseService,
|
||||
"get_by_id",
|
||||
lambda *_args, **_kwargs: (True, object()),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
sync_data_source.SyncLogsService,
|
||||
"increase_docs",
|
||||
lambda *_args, **_kwargs: None,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
sync_data_source.SyncLogsService,
|
||||
"duplicate_and_parse",
|
||||
lambda *_args, **_kwargs: (["parse error"], ["parsed-doc-id"]),
|
||||
)
|
||||
|
||||
task = {
|
||||
**_make_task(),
|
||||
"reindex": "0",
|
||||
"poll_range_start": datetime(2026, 1, 1, tzinfo=timezone.utc),
|
||||
"skip_connection_log": True,
|
||||
}
|
||||
sync = sync_data_source.BigQuery(_bigquery_conf(timestamp_column="updated_at"))
|
||||
|
||||
await sync._run_task_logic(task)
|
||||
|
||||
connector = _FakeBigQueryConnector.instance
|
||||
assert connector is not None
|
||||
assert connector.persist_sync_state_called is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.p2
|
||||
async def test_bigquery_cursor_does_not_persist_when_batch_is_skipped(monkeypatch):
|
||||
monkeypatch.setattr(sync_data_source, "BigQueryConnector", _FakeBigQueryConnector)
|
||||
_patch_common_dependencies(monkeypatch)
|
||||
monkeypatch.setattr(
|
||||
sync_data_source.KnowledgebaseService,
|
||||
"get_by_id",
|
||||
lambda *_args, **_kwargs: (True, object()),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
sync_data_source.SyncLogsService,
|
||||
"increase_docs",
|
||||
lambda *_args, **_kwargs: None,
|
||||
)
|
||||
|
||||
def _raise_in_duplicate_and_parse(*_args, **_kwargs):
|
||||
raise RuntimeError("batch failed")
|
||||
|
||||
monkeypatch.setattr(
|
||||
sync_data_source.SyncLogsService,
|
||||
"duplicate_and_parse",
|
||||
_raise_in_duplicate_and_parse,
|
||||
)
|
||||
|
||||
task = {
|
||||
**_make_task(),
|
||||
"reindex": "0",
|
||||
"poll_range_start": datetime(2026, 1, 1, tzinfo=timezone.utc),
|
||||
"skip_connection_log": True,
|
||||
}
|
||||
sync = sync_data_source.BigQuery(_bigquery_conf(timestamp_column="updated_at"))
|
||||
|
||||
await sync._run_task_logic(task)
|
||||
|
||||
connector = _FakeBigQueryConnector.instance
|
||||
assert connector is not None
|
||||
assert connector.persist_sync_state_called is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.p2
|
||||
async def test_bigquery_collect_prune_snapshot_when_enabled(monkeypatch):
|
||||
monkeypatch.setattr(sync_data_source, "BigQueryConnector", _FakeBigQueryConnector)
|
||||
|
||||
task = {
|
||||
**_make_task(),
|
||||
"reindex": "0",
|
||||
"poll_range_start": None,
|
||||
"skip_connection_log": True,
|
||||
}
|
||||
sync = sync_data_source.BigQuery(_bigquery_conf(sync_deleted_files=True))
|
||||
|
||||
await sync._generate(task)
|
||||
file_list = sync._collect_prune_snapshot(task)
|
||||
connector = _FakeBigQueryConnector.instance
|
||||
|
||||
assert connector.retrieve_all_slim_docs_perm_sync_called is True
|
||||
assert [doc.id for doc in file_list] == ["bq-row-1"]
|
||||
|
||||
|
||||
class _FakeDropboxConnector:
|
||||
instance = None
|
||||
|
||||
|
||||
Reference in New Issue
Block a user