#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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 logging
import os
import secrets
import time
import unittest
from dataclasses import dataclass
from typing import Any
from typing import Literal
from typing import Optional

import pytest
import sqlalchemy
from google.cloud.sql.connector import Connector
from parameterized import parameterized
from sqlalchemy import text

import apache_beam as beam
from apache_beam.io.jdbc import ReadFromJdbc
from apache_beam.ml.rag.ingestion import mysql_common
from apache_beam.ml.rag.ingestion import postgres_common
from apache_beam.ml.rag.ingestion import test_utils
from apache_beam.ml.rag.ingestion.base import VectorDatabaseWriteTransform
from apache_beam.ml.rag.ingestion.cloudsql import CloudSQLMySQLVectorWriterConfig
from apache_beam.ml.rag.ingestion.cloudsql import CloudSQLPostgresVectorWriterConfig
from apache_beam.ml.rag.ingestion.cloudsql import LanguageConnectorConfig
from apache_beam.ml.rag.types import Chunk
from apache_beam.ml.rag.types import Content
from apache_beam.ml.rag.types import Embedding
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to

_LOGGER = logging.getLogger(__name__)


@dataclass
class DatabaseTestConfig:
  """Database-specific test configuration."""
  database_type: Literal["postgresql", "mysql"]
  writer_config_class: type
  jdbc_driver: str
  connector_module: Literal["pg8000", "pymysql"]
  table_prefix: str

  password_env_var: str
  username: str
  database: str
  instance_uri: str

  vector_column_type: str
  metadata_column_type: str
  common_module: Any
  id_column_type: str = "VARCHAR(255)"


class DatabaseTestHelper:
  """Helper class to manage database setup, connections, and operations."""
  def __init__(self, db_config: DatabaseTestConfig, table_suffix: str):
    self.db_config = db_config
    self.table_suffix = table_suffix
    self.connector = None
    self.engine = None
    self.connection_config = None

    self.default_table_name = f"{db_config.table_prefix}{table_suffix}"
    self.custom_table_name = f"{db_config.table_prefix}_custom_{table_suffix}"
    self.metadata_conflicts_table = f"{db_config.table_prefix}_meta_conf_" \
      f"{table_suffix}"

    self._setup_read_queries()

  def _setup_read_queries(self):
    if self.db_config.database_type == "postgresql":
      self.read_queries = {
          self.default_table_name: f"""
                    SELECT 
                        CAST(id AS VARCHAR(255)),
                        CAST(content AS VARCHAR(255)),
                        CAST(embedding AS text),
                        CAST(metadata AS text)
                    FROM {self.default_table_name}
                """,
          self.custom_table_name: f"""
                    SELECT 
                        CAST(custom_id AS VARCHAR(255)),
                        CAST(embedding_vec AS text),
                        CAST(content_col AS VARCHAR(255)),
                        CAST(metadata AS text)
                    FROM {self.custom_table_name}
                    ORDER BY custom_id
                """,
          self.metadata_conflicts_table: f"""
                    SELECT 
                        CAST(id AS VARCHAR(255)),
                        CAST(embedding AS text),
                        CAST(content AS VARCHAR(255)),
                        CAST(source AS VARCHAR(255)),
                        CAST(timestamp AS VARCHAR(255))
                    FROM {self.metadata_conflicts_table}
                    ORDER BY timestamp, id
                """
      }
    elif self.db_config.database_type == "mysql":
      self.read_queries = {
          self.default_table_name: f"""
                    SELECT 
                        CAST(id AS CHAR(255)) as id,
                        CAST(content AS CHAR(255)) as content,
                        vector_to_string(embedding) as embedding,
                        CAST(metadata AS CHAR(10000)) as metadata
                    FROM {self.default_table_name}
                """,
          self.custom_table_name: f"""
                    SELECT 
                        CAST(custom_id AS CHAR(255)) as custom_id,
                        vector_to_string(embedding_vec) as embedding_vec,
                        CAST(content_col AS CHAR(255)) as content_col,
                        CAST(metadata AS CHAR(10000)) as metadata
                    FROM {self.custom_table_name}
                    ORDER BY custom_id
                """,
          self.metadata_conflicts_table: f"""
                    SELECT 
                        CAST(id AS CHAR(255)) as id,
                        vector_to_string(embedding) as embedding,
                        CAST(content AS CHAR(255)) as content,
                        CAST(source AS CHAR(255)) as source,
                        CAST(timestamp AS CHAR(255)) as timestamp
                    FROM {self.metadata_conflicts_table}
                    ORDER BY timestamp, id
                """
      }

  def get_read_query(self, table_name: str) -> str:
    if table_name not in self.read_queries:
      raise ValueError(f"No read query defined for table: {table_name}")
    return self.read_queries[table_name]

  def setup_connection(self):
    """Set up database connection and engine."""
    if not os.environ.get(self.db_config.password_env_var):
      raise ValueError("Password environment variable not set.")
    password = os.environ.get(self.db_config.password_env_var)

    self.connection_config = LanguageConnectorConfig(
        username=self.db_config.username,
        password=password,
        database_name=self.db_config.database,
        instance_name=self.db_config.instance_uri)

    self.connector = Connector(refresh_strategy="LAZY")

    def getconn():
      return self.connector.connect(
          self.db_config.instance_uri,
          self.db_config.connector_module,
          user=self.db_config.username,
          password=password,
          db=self.db_config.database,
      )

    dialect = "postgresql+pg8000" \
      if self.db_config.database_type == "postgresql" else "mysql+pymysql"
    self.engine = sqlalchemy.create_engine(f"{dialect}://", creator=getconn)

  def create_all_tables(self):
    if not self.engine:
      raise ValueError("Engine not initialized. Call setup_connection() first.")

    vector_type_large = self.db_config.vector_column_type.format(
        size=test_utils.VECTOR_SIZE)
    vector_type_small = self.db_config.vector_column_type.format(size=2)
    metadata_type = self.db_config.metadata_column_type
    id_type = self.db_config.id_column_type

    with self.engine.connect() as connection:
      default_table_sql = f"""
                CREATE TABLE {self.default_table_name} (
                    id {id_type} PRIMARY KEY,
                    embedding {vector_type_large},
                    content TEXT,
                    metadata {metadata_type}
                )
            """
      connection.execute(text(default_table_sql))

      custom_table_sql = f"""
                CREATE TABLE {self.custom_table_name} (
                    custom_id {id_type} PRIMARY KEY,
                    embedding_vec {vector_type_small},
                    content_col TEXT,
                    metadata {metadata_type}
                )
            """
      connection.execute(text(custom_table_sql))

      if self.db_config.database_type == "postgresql":
        metadata_conflicts_sql = f"""
                    CREATE TABLE {self.metadata_conflicts_table} (
                        id {id_type},
                        source TEXT,
                        timestamp TIMESTAMP,
                        content TEXT,
                        embedding {vector_type_small},
                        PRIMARY KEY (id),
                        UNIQUE (source, timestamp)
                    )
                """
      elif self.db_config.database_type == "mysql":
        metadata_conflicts_sql = f"""
                    CREATE TABLE {self.metadata_conflicts_table} (
                        id {id_type},
                        source TEXT,
                        timestamp TIMESTAMP,
                        content TEXT,
                        embedding {vector_type_small},
                        PRIMARY KEY (id),
                        UNIQUE KEY unique_source_timestamp (source(255), timestamp)
                    )
                """
      connection.execute(text(metadata_conflicts_sql))
      connection.commit()

  def create_writer_config(
      self,
      table_name: Optional[str] = None,
      column_specs=None,
      conflict_resolution=None):
    if not self.connection_config:
      raise ValueError(
          "Connection not initialized. Call setup_connection() first.")

    table_name = table_name or self.default_table_name

    kwargs = {
        'connection_config': self.connection_config,
        'table_name': table_name,
    }

    if column_specs is not None:
      kwargs['column_specs'] = column_specs
    if conflict_resolution is not None:
      kwargs['conflict_resolution'] = conflict_resolution

    return self.db_config.writer_config_class(**kwargs)

  def cleanup(self):
    if self.engine:
      table_names = [
          self.default_table_name,
          self.custom_table_name,
          self.metadata_conflicts_table
      ]

      try:
        with self.engine.connect() as connection:
          for table_name in table_names:
            connection.execute(text(f"DROP TABLE IF EXISTS {table_name}"))
          connection.commit()
        _LOGGER.info(
            "Dropped %s tables: %s",
            self.db_config.database_type,
            ', '.join(table_names))
      except Exception as e:
        _LOGGER.warning(
            "Error dropping %s tables: %s", self.db_config.database_type, e)

    if self.connector:
      try:
        self.connector.close()
      except Exception as e:
        _LOGGER.warning("Error closing connector: %s", e)

    if self.engine:
      try:
        self.engine.dispose()
      except Exception as e:
        _LOGGER.warning("Error disposing engine: %s", e)


class PipelineVerificationHelper:
  """Helper class for common pipeline verification patterns."""
  @staticmethod
  def build_jdbc_params(helper: DatabaseTestHelper, table_name: str) -> dict:
    """Build JDBC parameters dictionary for ReadFromJdbc."""
    writer_config = helper.create_writer_config(table_name)

    return {
        'table_name': table_name,
        'driver_class_name': helper.db_config.jdbc_driver,
        'jdbc_url': writer_config.connector_config.to_connection_config().
        jdbc_url,
        'username': helper.db_config.username,
        'password': helper.connection_config.password,
        'query': helper.get_read_query(table_name),
        'classpath': writer_config.connector_config.additional_jdbc_args()
        ['classpath']
    }

  @staticmethod
  def verify_standard_operations(
      pipeline, jdbc_params: dict, expected_chunks: list[Chunk]):
    num_records = len(expected_chunks)
    sample_size = min(500, num_records // 2)

    with pipeline as p:
      rows = (p | ReadFromJdbc(**jdbc_params))

      # Count verification
      count_result = rows | "Count All" >> beam.combiners.Count.Globally()
      assert_that(count_result, equal_to([num_records]), label='count_check')

      # Hash verification
      chunks = (rows | "To Chunks" >> beam.Map(test_utils.row_to_chunk))
      chunk_hashes = chunks | "Hash Chunks" >> beam.CombineGlobally(
          test_utils.HashingFn())
      expected_hash = test_utils.generate_expected_hash(num_records)
      assert_that(chunk_hashes, equal_to([expected_hash]), label='hash_check')

      # Sample validation - first N
      first_n = (
          chunks
          | "Key on Index" >> beam.Map(test_utils.key_on_id)
          | f"Get First {sample_size}" >> beam.transforms.combiners.Top.Of(
              sample_size, key=lambda x: x[0], reverse=True)
          | "Remove Keys 1" >> beam.Map(lambda xs: [x[1] for x in xs]))
      expected_first_n = expected_chunks[:sample_size]
      assert_that(
          first_n,
          equal_to([expected_first_n]),
          label=f"first_{sample_size}_check")

      # Sample validation - last N
      last_n = (
          chunks
          | "Key on Index 2" >> beam.Map(test_utils.key_on_id)
          | f"Get Last {sample_size}" >> beam.transforms.combiners.Top.Of(
              sample_size, key=lambda x: x[0])
          | "Remove Keys 2" >> beam.Map(lambda xs: [x[1] for x in xs]))
      expected_last_n = expected_chunks[-sample_size:][::-1]
      assert_that(
          last_n,
          equal_to([expected_last_n]),
          label=f"last_{sample_size}_check")


# Database configurations
POSTGRES_CONFIG = DatabaseTestConfig(
    database_type="postgresql",
    writer_config_class=CloudSQLPostgresVectorWriterConfig,
    jdbc_driver="org.postgresql.Driver",
    connector_module="pg8000",
    table_prefix="python_rag_postgres_",
    password_env_var="ALLOYDB_PASSWORD",
    username="postgres",
    database="postgres",
    instance_uri="apache-beam-testing:us-central1:beam-integration-tests",
    vector_column_type="VECTOR({size})",
    metadata_column_type="JSONB",
    common_module=postgres_common)

MYSQL_CONFIG = DatabaseTestConfig(
    database_type="mysql",
    writer_config_class=CloudSQLMySQLVectorWriterConfig,
    jdbc_driver="com.mysql.cj.jdbc.Driver",
    connector_module="pymysql",
    table_prefix="python_rag_mysql_",
    password_env_var="ALLOYDB_PASSWORD",
    username="mysql",
    database="embeddings",
    instance_uri="apache-beam-testing:us-central1:beam-integration-tests-mysql",
    vector_column_type="VECTOR({size}) USING VARBINARY",
    metadata_column_type="JSON",
    common_module=mysql_common)


@pytest.mark.uses_gcp_java_expansion_service
@unittest.skipUnless(
    os.environ.get('EXPANSION_JARS'),
    "EXPANSION_JARS environment var is not provided, "
    "indicating that jars have not been built")
class CloudSQLVectorWriterConfigTest(unittest.TestCase):
  def setUp(self):
    self.write_test_pipeline = TestPipeline(is_integration_test=True)
    self.read_test_pipeline = TestPipeline(is_integration_test=True)
    self.write_test_pipeline2 = TestPipeline(is_integration_test=True)
    self.read_test_pipeline2 = TestPipeline(is_integration_test=True)

    self.write_test_pipeline.not_use_test_runner_api = True
    self.read_test_pipeline.not_use_test_runner_api = True
    self.write_test_pipeline2.not_use_test_runner_api = True
    self.read_test_pipeline2.not_use_test_runner_api = True
    self._runner = type(self.read_test_pipeline.runner).__name__

    self.db_helpers = {}
    self.table_suffix = '%d%s' % (int(time.time()), secrets.token_hex(3))

    # Set up database helpers
    for config in [POSTGRES_CONFIG, MYSQL_CONFIG]:
      helper = DatabaseTestHelper(config, self.table_suffix)
      helper.setup_connection()
      helper.create_all_tables()
      self.db_helpers[config.database_type] = helper
      _LOGGER.info("Successfully set up %s database", config.database_type)

  def tearDown(self):
    for helper in self.db_helpers.values():
      helper.cleanup()

  def skip_if_dataflow_runner(self):
    if self._runner and "dataflowrunner" in self._runner.lower():
      self.skipTest(
          "Skipping some tests on Dataflow Runner to avoid bloat and timeouts")

  @parameterized.expand([(POSTGRES_CONFIG), (MYSQL_CONFIG)])
  def test_default_config(self, db_config):
    """Test basic write and read operations with default configuration.
      
      This test validates the most basic CloudSQL vector database functionality:
      - Default table schema: id (VARCHAR), content (TEXT), embedding (VECTOR),
        metadata (JSON/JSONB) 
      - Default column specifications (no customization)
      - Default conflict resolution (IGNORE on primary key conflicts)
      - Write chunks to database and read them back
      - Verify data integrity through count, hash, and sample validation
      """
    self.skip_if_dataflow_runner()

    helper = self.db_helpers[db_config.database_type]
    num_records = 150

    # Create test data
    test_chunks = test_utils.ChunkTestUtils.get_expected_values(0, num_records)

    # Write test
    writer_config = helper.create_writer_config()
    self.write_test_pipeline.not_use_test_runner_api = True
    with self.write_test_pipeline as p:
      _ = (
          p | beam.Create(test_chunks)
          | VectorDatabaseWriteTransform(writer_config))

    # Read and verify
    self.read_test_pipeline.not_use_test_runner_api = True
    jdbc_params = PipelineVerificationHelper.build_jdbc_params(
        helper, helper.default_table_name)
    PipelineVerificationHelper.verify_standard_operations(
        self.read_test_pipeline, jdbc_params, test_chunks)

  @parameterized.expand([
      (POSTGRES_CONFIG, "UPDATE", ["embedding", "content"]),
      (MYSQL_CONFIG, "UPDATE", ["embedding", "content"]),
      (POSTGRES_CONFIG, "IGNORE", None),
      (MYSQL_CONFIG, "IGNORE", None),
      (POSTGRES_CONFIG, "UPDATE_ALL", None),  # Default update fields
      (MYSQL_CONFIG, "UPDATE_ALL", None),
  ])
  def test_conflict_resolution(self, db_config, action, update_fields):
    """Test conflict resolution strategies when primary key conflicts occur.
    
      This test validates different approaches to handling duplicate primary
      keys:
      
      UPDATE with specific fields:
      - When duplicate ID encountered, update only specified fields (embedding,
        content)
      - Other fields (metadata) remain unchanged from original record
      
      IGNORE:
      - When duplicate ID encountered, keep original record unchanged
      
      UPDATE_ALL (default update fields):
      - When duplicate ID encountered, update ALL non-key fields
      - This includes content, embedding, AND metadata
      
      Scenario for all strategies:
      1. Insert initial records
      2. Insert records with same IDs but different content/embeddings  
      3. Verify final state matches expected conflict resolution behavior
      """
    self.skip_if_dataflow_runner()

    helper = self.db_helpers[db_config.database_type]
    num_records = 20

    common_module = db_config.common_module
    if action == "IGNORE":
      if db_config.database_type == "mysql":
        conflict_resolution = common_module.ConflictResolution(
            action="IGNORE", primary_key_field="id")
      else:
        conflict_resolution = None  # Default behavior for PostgreSQL
    elif action == "UPDATE":
      if db_config.database_type == "postgresql":
        conflict_resolution = common_module.ConflictResolution(
            on_conflict_fields="id",
            action="UPDATE",
            update_fields=update_fields)
      else:
        conflict_resolution = common_module.ConflictResolution(
            action="UPDATE", update_fields=update_fields)
    else:  # UPDATE_ALL
      if db_config.database_type == "postgresql":
        conflict_resolution = common_module.ConflictResolution(
            on_conflict_fields="id", action="UPDATE")
      else:
        conflict_resolution = common_module.ConflictResolution(action="UPDATE")

    initial_chunks = test_utils.ChunkTestUtils.get_expected_values(
        0, num_records)
    writer_config = helper.create_writer_config(
        conflict_resolution=conflict_resolution)

    self.write_test_pipeline.not_use_test_runner_api = True
    with self.write_test_pipeline as p:
      _ = (
          p | "Write Initial" >> beam.Create(initial_chunks)
          | VectorDatabaseWriteTransform(writer_config))

    # Write conflicting data
    updated_chunks = test_utils.ChunkTestUtils.get_expected_values(
        0, num_records, content_prefix="Updated", seed_multiplier=2)

    self.write_test_pipeline2.not_use_test_runner_api = True
    with self.write_test_pipeline2 as p:
      _ = (
          p | "Write Conflicts" >> beam.Create(updated_chunks)
          | VectorDatabaseWriteTransform(writer_config))

    jdbc_params = PipelineVerificationHelper.build_jdbc_params(
        helper, helper.default_table_name)
    expected_chunks = updated_chunks if action != "IGNORE" else initial_chunks

    self.read_test_pipeline.not_use_test_runner_api = True
    with self.read_test_pipeline as p:
      rows = (p | ReadFromJdbc(**jdbc_params))

      count_result = rows | "Count All" >> beam.combiners.Count.Globally()
      assert_that(count_result, equal_to([num_records]), label='count_check')

      chunks = rows | "To Chunks" >> beam.Map(test_utils.row_to_chunk)
      assert_that(chunks, equal_to(expected_chunks), label='chunks_check')

  @parameterized.expand([(POSTGRES_CONFIG), (MYSQL_CONFIG)])
  def test_custom_column_names_and_value_functions(self, db_config):
    """Test completely custom column specifications with custom value
      extraction.
    
      This test validates advanced customization of how chunk data is stored:
      
      Custom column names:
      - custom_id (instead of 'id')
      - embedding_vec (instead of 'embedding') 
      - content_col (instead of 'content')
      
      Custom value extraction functions:
      - ID: Extract timestamp from metadata and prefix with "timestamp_"
      - Content: Prefix content with its character length "10:actual_content"
      - Embedding: Use custom embedding extraction function
      
      This tests the flexibility to completely reshape how chunk data maps 
      to database columns, useful for integrating with existing database schemas
      or applying business-specific transformations.
      """
    self.skip_if_dataflow_runner()

    helper = self.db_helpers[db_config.database_type]
    num_records = 20
    common_module = db_config.common_module

    test_chunks = [
        Chunk(
            id=str(i),
            content=Content(text=f"content_{i}"),
            embedding=Embedding(dense_embedding=[float(i), float(i + 1)]),
            metadata={"timestamp": f"2024-02-02T{i:02d}:00:00"})
        for i in range(num_records)
    ]

    chunk_embedding_fn = common_module.chunk_embedding_fn
    specs = (
        common_module.ColumnSpecsBuilder().add_custom_column_spec(
            common_module.ColumnSpec.text(
                column_name="custom_id",
                value_fn=lambda chunk:
                f"timestamp_{chunk.metadata.get('timestamp', '')}")
        ).add_custom_column_spec(
            common_module.ColumnSpec.vector(
                column_name="embedding_vec",
                value_fn=chunk_embedding_fn)).add_custom_column_spec(
                    common_module.ColumnSpec.text(
                        column_name="content_col",
                        value_fn=lambda chunk:
                        f"{len(chunk.content.text)}:{chunk.content.text}")).
        with_metadata_spec().build())

    def custom_row_to_chunk(row):
      timestamp = row.custom_id.split('timestamp_')[1]
      i = int(timestamp.split('T')[1][:2])

      embedding_list = [
          float(x) for x in row.embedding_vec.strip('[]').split(',')
      ]

      content = row.content_col.split(':', 1)[1]

      return Chunk(
          id=str(i),
          content=Content(text=content),
          embedding=Embedding(dense_embedding=embedding_list),
          metadata=json.loads(row.metadata))

    writer_config = helper.create_writer_config(helper.custom_table_name, specs)
    self.write_test_pipeline.not_use_test_runner_api = True
    with self.write_test_pipeline as p:
      _ = (
          p | beam.Create(test_chunks)
          | VectorDatabaseWriteTransform(writer_config))

    jdbc_params = PipelineVerificationHelper.build_jdbc_params(
        helper, helper.custom_table_name)

    self.read_test_pipeline.not_use_test_runner_api = True
    with self.read_test_pipeline as p:
      rows = (p | ReadFromJdbc(**jdbc_params))

      count_result = rows | "Count All" >> beam.combiners.Count.Globally()
      assert_that(count_result, equal_to([num_records]), label='count_check')

      chunks = rows | "To Chunks" >> beam.Map(custom_row_to_chunk)
      assert_that(chunks, equal_to(test_chunks), label='chunks_check')

  @parameterized.expand([(POSTGRES_CONFIG), (MYSQL_CONFIG)])
  def test_custom_type_conversion_with_default_columns(self, db_config):
    """Test custom type conversion and SQL typecasting with modified column
      names.
    
      This test validates data type handling and database-specific SQL features:
      
      Type conversion:
      - Convert string IDs to integers before storage
      - Apply length-prefix transformation to content
      
      SQL typecasting (database-specific):
      - PostgreSQL: Use ::text typecast for converted integers
      - MySQL: Rely on automatic type conversion (no explicit typecast)
      
      Column name customization:
      - Use custom names but with standard spec builders (not completely custom
        functions)
      
      This tests the ability to adapt data types for database constraints
      while maintaining the standard chunk-to-database mapping logic.
    """
    self.skip_if_dataflow_runner()

    helper = self.db_helpers[db_config.database_type]
    num_records = 20
    common_module = db_config.common_module

    test_chunks = [
        Chunk(
            id=str(i),
            content=Content(text=f"content_{i}"),
            embedding=Embedding(dense_embedding=[float(i), float(i + 1)]),
            metadata={"timestamp": f"2024-02-02T{i:02d}:00:00"})
        for i in range(num_records)
    ]

    if db_config.database_type == "postgresql":
      specs = (
          common_module.ColumnSpecsBuilder().with_id_spec(
              column_name="custom_id",
              python_type=int,
              convert_fn=lambda x: int(x),
              sql_typecast="::text").with_content_spec(
                  column_name="content_col",
                  convert_fn=lambda x: f"{len(x)}:{x}"  # Add length prefix
              ).with_embedding_spec(
                  column_name="embedding_vec").with_metadata_spec().build())
    else:  # MySQL
      specs = (
          common_module.ColumnSpecsBuilder().with_id_spec(
              column_name="custom_id",
              python_type=int,
              convert_fn=lambda x: int(x)).with_content_spec(
                  column_name="content_col",
                  convert_fn=lambda x: f"{len(x)}:{x}").with_embedding_spec(
                      column_name="embedding_vec").with_metadata_spec().build())

    def type_conversion_row_to_chunk(row):
      embedding_list = [
          float(x) for x in row.embedding_vec.strip('[]').split(',')
      ]

      content = row.content_col.split(':', 1)[1]

      return Chunk(
          id=row.custom_id,  # custom_id is the converted ID field
          content=Content(text=content),
          embedding=Embedding(dense_embedding=embedding_list),
          metadata=json.loads(row.metadata))

    writer_config = helper.create_writer_config(helper.custom_table_name, specs)
    self.write_test_pipeline.not_use_test_runner_api = True
    with self.write_test_pipeline as p:
      _ = (
          p | beam.Create(test_chunks)
          | VectorDatabaseWriteTransform(writer_config))

    jdbc_params = PipelineVerificationHelper.build_jdbc_params(
        helper, helper.custom_table_name)

    self.read_test_pipeline.not_use_test_runner_api = True
    with self.read_test_pipeline as p:
      rows = (p | ReadFromJdbc(**jdbc_params))

      count_result = rows | "Count All" >> beam.combiners.Count.Globally()
      assert_that(count_result, equal_to([num_records]), label='count_check')

      chunks = rows | "To Chunks" >> beam.Map(type_conversion_row_to_chunk)
      assert_that(chunks, equal_to(test_chunks), label='chunks_check')

  @parameterized.expand([(POSTGRES_CONFIG), (MYSQL_CONFIG)])
  def test_default_id_embedding_specs(self, db_config):
    """Test minimal schema with only ID and embedding columns.
  
      This test validates the ability to create a minimal vector database
      schema:
      - Only stores id and embedding fields
      - content and metadata columns are excluded from the table
      - Tests that the system correctly handles missing/null fields
      
      Use case: When you only need vector similarity search without storing 
      the original content or metadata (perhaps stored elsewhere).
      
      Validation:
      - Chunks written with content/metadata are stored with those fields as
        null
      - Reading back produces chunks with null content and empty metadata
      - Vector similarity operations still work normally
      """
    self.skip_if_dataflow_runner()

    helper = self.db_helpers[db_config.database_type]
    num_records = 20
    common_module = db_config.common_module

    specs = (
        common_module.ColumnSpecsBuilder().with_id_spec().with_embedding_spec().
        build())

    writer_config = helper.create_writer_config(column_specs=specs)

    test_chunks = test_utils.ChunkTestUtils.get_expected_values(0, num_records)

    with self.write_test_pipeline as p:
      _ = (
          p | beam.Create(test_chunks)
          | VectorDatabaseWriteTransform(writer_config))

    expected_chunks = test_utils.ChunkTestUtils.get_expected_values(
        0, num_records)
    for chunk in expected_chunks:
      chunk.content.text = None  # Content column not included in schema
      chunk.metadata = {}  # Metadata column not included in schema

    jdbc_params = PipelineVerificationHelper.build_jdbc_params(
        helper, helper.default_table_name)
    if db_config.database_type == "postgresql":
      jdbc_params['query'] = f"""
            SELECT 
                CAST(id AS VARCHAR(255)),
                CAST(embedding AS text)
            FROM {helper.default_table_name}
            ORDER BY id
        """
    elif db_config.database_type == "mysql":
      jdbc_params['query'] = f"""
            SELECT 
                CAST(id AS CHAR(255)) as id,
                vector_to_string(embedding) as embedding
            FROM {helper.default_table_name}
        """

    with self.read_test_pipeline as p:
      rows = (p | ReadFromJdbc(**jdbc_params))
      chunks = rows | "To Chunks" >> beam.Map(test_utils.row_to_chunk)
      assert_that(chunks, equal_to(expected_chunks), label='chunks_check')

  @parameterized.expand([(POSTGRES_CONFIG), (MYSQL_CONFIG)])
  def test_metadata_field_extraction(self, db_config):
    """Test extracting specific metadata fields into separate database columns.
      
      This test validates the ability to:
      - Extract specific fields from the JSON metadata object 
      - Map them to dedicated database columns (e.g., metadata.source -> source
        column)
      - Apply database-specific SQL typecasts (PostgreSQL ::timestamp vs MySQL
        default)
      - Store and retrieve the extracted fields correctly
      
      This is different from default metadata handling which stores the entire 
      metadata object as JSON in a single column.
      """
    self.skip_if_dataflow_runner()

    helper = self.db_helpers[db_config.database_type]
    num_records = 20
    common_module = db_config.common_module

    if db_config.database_type == "postgresql":
      specs = (
          common_module.ColumnSpecsBuilder().with_id_spec().with_embedding_spec(
          ).with_content_spec().add_metadata_field(
              field="source",
              column_name="source",
              python_type=str,
              sql_typecast=None).add_metadata_field(
                  field="timestamp",
                  python_type=str,
                  sql_typecast="::timestamp").build())
    else:
      specs = (
          common_module.ColumnSpecsBuilder().with_id_spec().with_embedding_spec(
          ).with_content_spec().add_metadata_field(
              field="source", column_name="source",
              python_type=str).add_metadata_field(
                  field="timestamp", python_type=str).build())

    writer_config = helper.create_writer_config(
        helper.metadata_conflicts_table, specs, conflict_resolution=None)

    test_chunks = [
        Chunk(
            id=str(i),
            content=Content(text=f"content_{i}"),
            embedding=Embedding(dense_embedding=[float(i), float(i + 1)]),
            metadata={
                "source": f"source_{i % 3}",
                "timestamp": f"2024-02-02T{i:02d}:00:00"
            }) for i in range(num_records)
    ]

    self.write_test_pipeline.not_use_test_runner_api = True
    with self.write_test_pipeline as p:
      _ = (
          p | beam.Create(test_chunks)
          | VectorDatabaseWriteTransform(writer_config))

    def metadata_row_to_chunk(row):
      embedding_list = [float(x) for x in row.embedding.strip('[]').split(',')]
      timestamp = row.timestamp.replace(
          ' ', 'T') if ' ' in row.timestamp else row.timestamp
      return Chunk(
          id=row.id,
          content=Content(text=row.content),
          embedding=Embedding(dense_embedding=embedding_list),
          metadata={
              "source": row.source, "timestamp": timestamp
          })

    jdbc_params = PipelineVerificationHelper.build_jdbc_params(
        helper, helper.metadata_conflicts_table)

    self.read_test_pipeline.not_use_test_runner_api = True
    with self.read_test_pipeline as p:
      rows = (p | ReadFromJdbc(**jdbc_params))
      chunks = rows | "To Chunks" >> beam.Map(metadata_row_to_chunk)
      assert_that(chunks, equal_to(test_chunks), label='chunks_check')

  @parameterized.expand([(POSTGRES_CONFIG), (MYSQL_CONFIG)])
  def test_composite_unique_constraint_conflicts(self, db_config):
    """Test conflict resolution when unique constraints span multiple columns.
      
      This test validates conflict resolution when the unique constraint is NOT 
      on the primary key, but on a combination of other columns (source +
      timestamp).
      
      Scenario:
      1. Insert records with unique (source, timestamp) combinations
      2. Attempt to insert records with same (source, timestamp) but different
         IDs and content
      3. Verify that conflict resolution (UPDATE) works correctly based on
         composite key
      
      This is different from test_conflict_resolution which tests conflicts on 
      the primary key field only.
      """
    self.skip_if_dataflow_runner()

    helper = self.db_helpers[db_config.database_type]
    num_records = 5
    common_module = db_config.common_module

    if db_config.database_type == "postgresql":
      specs = (
          common_module.ColumnSpecsBuilder().with_id_spec().with_embedding_spec(
          ).with_content_spec().add_metadata_field(
              field="source",
              column_name="source",
              python_type=str,
              sql_typecast=None).add_metadata_field(
                  field="timestamp",
                  python_type=str,
                  sql_typecast="::timestamp").build())

      conflict_resolution = common_module.ConflictResolution(
          on_conflict_fields=["source", "timestamp"],
          action="UPDATE",
          update_fields=["embedding", "content"])
    elif db_config.database_type == "mysql":
      specs = (
          common_module.ColumnSpecsBuilder().with_id_spec().with_embedding_spec(
          ).with_content_spec().add_metadata_field(
              field="source", column_name="source",
              python_type=str).add_metadata_field(
                  field="timestamp", python_type=str).build())

      # MySQL conflict resolution - detects unique constraint automatically
      conflict_resolution = common_module.ConflictResolution(
          action="UPDATE", update_fields=["embedding", "content"])

    writer_config = helper.create_writer_config(
        helper.metadata_conflicts_table, specs, conflict_resolution)

    initial_chunks = [
        Chunk(
            id=str(i),
            content=Content(text=f"content_{i}"),
            embedding=Embedding(dense_embedding=[float(i), float(i + 1)]),
            metadata={
                "source": "source_A", "timestamp": f"2024-02-02T{i:02d}:00:00"
            }) for i in range(num_records)
    ]

    with self.write_test_pipeline as p:
      _ = (
          p | "Write Initial" >> beam.Create(initial_chunks)
          | VectorDatabaseWriteTransform(writer_config))

    conflicting_chunks = [
        Chunk(
            id=f"new_{i}",
            content=Content(text=f"updated_content_{i}"),
            embedding=Embedding(
                dense_embedding=[float(i) * 2, float(i + 1) * 2]),
            metadata={
                "source": "source_A", "timestamp": f"2024-02-02T{i:02d}:00:00"
            }) for i in range(num_records)
    ]

    with self.write_test_pipeline2 as p:
      _ = (
          p | "Write Conflicts" >> beam.Create(conflicting_chunks)
          | VectorDatabaseWriteTransform(writer_config))

    expected_chunks = [
        Chunk(
            id=str(i),
            content=Content(text=f"updated_content_{i}"),
            embedding=Embedding(
                dense_embedding=[float(i) * 2, float(i + 1) * 2]),
            metadata={
                "source": "source_A", "timestamp": f"2024-02-02T{i:02d}:00:00"
            }) for i in range(num_records)
    ]

    def metadata_row_to_chunk(row):
      embedding_list = [float(x) for x in row.embedding.strip('[]').split(',')]
      timestamp = row.timestamp.replace(
          ' ', 'T') if ' ' in row.timestamp else row.timestamp
      return Chunk(
          id=row.id,
          content=Content(text=row.content),
          embedding=Embedding(dense_embedding=embedding_list),
          metadata={
              "source": row.source, "timestamp": timestamp
          })

    jdbc_params = PipelineVerificationHelper.build_jdbc_params(
        helper, helper.metadata_conflicts_table)

    with self.read_test_pipeline as p:
      rows = (p | ReadFromJdbc(**jdbc_params))

      count_result = rows | "Count All" >> beam.combiners.Count.Globally()
      assert_that(count_result, equal_to([num_records]), label='count_check')

      chunks = rows | "To Chunks" >> beam.Map(metadata_row_to_chunk)
      assert_that(chunks, equal_to(expected_chunks), label='chunks_check')


if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  unittest.main()
