sdks/python/apache_beam/ml/rag/ingestion/cloudsql.py - beam - Git at Google


            
            #

            
            # 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.


            from dataclasses import asdict

            
            from dataclasses import dataclass

            
            from typing import Any

            
            from typing import Optional


            from apache_beam.ml.rag.ingestion import mysql

            
            from apache_beam.ml.rag.ingestion import mysql_common

            
            from apache_beam.ml.rag.ingestion import postgres

            
            from apache_beam.ml.rag.ingestion import postgres_common

            
            from apache_beam.ml.rag.ingestion.jdbc_common import ConnectionConfig

            
            from apache_beam.ml.rag.ingestion.jdbc_common import WriteConfig


            @dataclass

            
            class LanguageConnectorConfig:

            
              """Configuration options for CloudSQL Java language connector.


                Set parameters to connect connection to a CloudSQL instance using

            
                Java language connector connector. For details see

            
                https://github.com/GoogleCloudPlatform/cloud-sql-jdbc-socket-factory/blob/main/docs/jdbc.md


                Attributes:

            
                    username: Database username.

            
                    password: Database password. Can be empty string when using IAM.

            
                    database_name: Name of the database to connect to.

            
                    instance_name: Instance connection name. Format:

            
                        '<PROJECT>:<REGION>:<INSTANCE>'

            
                    ip_type: Preferred order of IP types used to connect via a comma

            
                        list of strings.

            
                    enable_iam_auth: Whether to enable IAM authentication. Default is False

            
                    target_principal: Optional service account to impersonate for

            
                        connection.

            
                    delegates: Optional list of service accounts for delegated

            
                        impersonation.

            
                    admin_service_endpoint: Optional custom API service endpoint.

            
                    quota_project: Optional project ID for quota and billing.

            
                    connection_properties: Optional JDBC connection properties dict.

            
                        Example: {'ssl': 'true'}

            
                    additional_properties: Additional properties to be added to the JDBC

            
                        url. Example: {'someProperty': 'true'}

            
                """

            
              username: str

            
              password: str

            
              database_name: str

            
              instance_name: str

            
              ip_types: Optional[list[str]] = None

            
              enable_iam_auth: bool = False

            
              target_principal: Optional[str] = None

            
              delegates: Optional[list[str]] = None

            
              quota_project: Optional[str] = None

            
              connection_properties: Optional[dict[str, str]] = None

            
              additional_properties: Optional[dict[str, Any]] = None


              def _base_jdbc_properties(self) -> dict[str, Any]:

            
                properties = {"cloudSqlInstance": self.instance_name}


                if self.ip_types:

            
                  properties["ipTypes"] = ",".join(self.ip_types)


                if self.enable_iam_auth:

            
                  properties["enableIamAuth"] = "true"


                if self.target_principal:

            
                  properties["cloudSqlTargetPrincipal"] = self.target_principal


                if self.delegates:

            
                  properties["cloudSqlDelegates"] = ",".join(self.delegates)


                if self.quota_project:

            
                  properties["cloudSqlAdminQuotaProject"] = self.quota_project


                if self.additional_properties:

            
                  properties.update(self.additional_properties)


                return properties


              def _build_jdbc_url(self, socketFactory, database_type):

            
                url = f"jdbc:{database_type}:///{self.database_name}?"


                properties = self._base_jdbc_properties()

            
                properties['socketFactory'] = socketFactory


                property_string = "&".join(f"{k}={v}" for k, v in properties.items())

            
                return url + property_string


              def to_connection_config(self):

            
                return ConnectionConfig(

            
                    jdbc_url=self.to_jdbc_url(),

            
                    username=self.username,

            
                    password=self.password,

            
                    connection_properties=self.connection_properties,

            
                    additional_jdbc_args=self.additional_jdbc_args())


              def additional_jdbc_args(self) -> dict[str, list[Any]]:

            
                return {}


            @dataclass

            
            class _PostgresConnectorConfig(LanguageConnectorConfig):

            
              def to_jdbc_url(self) -> str:

            
                """Convert options to a properly formatted JDBC URL.


                  Returns:

            
                      JDBC URL string configured with all options.

            
                  """

            
                return self._build_jdbc_url(

            
                    socketFactory="com.google.cloud.sql.postgres.SocketFactory",

            
                    database_type="postgresql")


              def additional_jdbc_args(self) -> dict[str, list[Any]]:

            
                return {

            
                    'classpath': [

            
                        "org.postgresql:postgresql:42.2.16",

            
                        "com.google.cloud.sql:postgres-socket-factory:1.25.0"

            
                    ]

            
                }


              @classmethod

            
              def from_base_config(cls, config: LanguageConnectorConfig):

            
                return cls(**asdict(config))


            class CloudSQLPostgresVectorWriterConfig(postgres.PostgresVectorWriterConfig):

            
              def __init__(

            
                  self,

            
                  connection_config: LanguageConnectorConfig,

            
                  table_name: str,

            
                  *,

            
                  # pylint: disable=dangerous-default-value

            
                  write_config: WriteConfig = WriteConfig(),

            
                  column_specs: list[postgres_common.ColumnSpec] = postgres_common.

            
                  ColumnSpecsBuilder.with_defaults().build(),

            
                  conflict_resolution: Optional[

            
                      postgres_common.ConflictResolution] = postgres_common.

            
                  ConflictResolution(on_conflict_fields=[], action='IGNORE')):

            
                """Configuration for writing vectors to ClouSQL Postgres.


                Supports flexible schema configuration through column specifications and

            
                conflict resolution strategies.


                Args:

            
                    connection_config: :class:`LanguageConnectorConfig`.

            
                    table_name: Target table name.

            
                    write_config: JdbcIO :class:`~.jdbc_common.WriteConfig` to control

            
                      batch sizes, authosharding, etc.

            
                    column_specs:

            
                        Use :class:`~.postgres_common.ColumnSpecsBuilder` to configure how

            
                        embeddings and metadata are written a database

            
                        schema. If None, uses default EmbeddableItem schema.

            
                    conflict_resolution: Optional

            
                        :class:`~.postgres_common.ConflictResolution`

            
                        strategy for handling insert conflicts. ON CONFLICT DO NOTHING by

            
                        default.


                Examples:

            
                    Basic usage with default schema:


                    >>> config = PostgresVectorWriterConfig(

            
                    ...     connection_config=PostgresConnectionConfig(...),

            
                    ...     table_name='embeddings'

            
                    ... )


                    Simple case with default schema:


                    >>> config = PostgresVectorWriterConfig(

            
                    ...     connection_config=ConnectionConfig(...),

            
                    ...     table_name='embeddings'

            
                    ... )


                    Custom schema with metadata fields:


                    >>> specs = (ColumnSpecsBuilder()

            
                    ...         .with_id_spec(column_name="my_id_column")

            
                    ...         .with_embedding_spec(column_name="embedding_vec")

            
                    ...         .add_metadata_field(field="source", column_name="src")

            
                    ...         .add_metadata_field(

            
                    ...             "timestamp",

            
                    ...             column_name="created_at",

            
                    ...             sql_typecast="::timestamp"

            
                    ...         )

            
                    ...         .build())


                    Minimal schema (only ID + embedding written)


                    >>> column_specs = (ColumnSpecsBuilder()

            
                    ...     .with_id_spec()

            
                    ...     .with_embedding_spec()

            
                    ...     .build())


                    >>> config = CloudSQLPostgresVectorWriterConfig(

            
                    ...     connection_config=PostgresConnectionConfig(...),

            
                    ...     table_name='embeddings',

            
                    ...     column_specs=specs

            
                    ... )

            
                """

            
                self.connector_config = _PostgresConnectorConfig.from_base_config(

            
                    connection_config)

            
                super().__init__(

            
                    connection_config=self.connector_config.to_connection_config(),

            
                    write_config=write_config,

            
                    table_name=table_name,

            
                    column_specs=column_specs,

            
                    conflict_resolution=conflict_resolution)


            @dataclass

            
            class _MySQLConnectorConfig(LanguageConnectorConfig):

            
              def to_jdbc_url(self) -> str:

            
                """Convert options to a properly formatted MySQL JDBC URL."""

            
                return self._build_jdbc_url(

            
                    socketFactory="com.google.cloud.sql.mysql.SocketFactory",

            
                    database_type="mysql")


              def additional_jdbc_args(self) -> dict[str, list[Any]]:

            
                return {

            
                    'classpath': [

            
                        "mysql:mysql-connector-java:8.0.22",

            
                        "com.google.cloud.sql:mysql-socket-factory-connector-j-8:1.25.0"

            
                    ]

            
                }


              @classmethod

            
              def from_base_config(cls, config: LanguageConnectorConfig):

            
                return cls(**asdict(config))


            class CloudSQLMySQLVectorWriterConfig(mysql.MySQLVectorWriterConfig):

            
              def __init__(

            
                  self,

            
                  connection_config: LanguageConnectorConfig,

            
                  table_name: str,

            
                  *,

            
                  write_config: WriteConfig = WriteConfig(),

            
                  # pylint: disable=dangerous-default-value

            
                  column_specs: list[mysql_common.ColumnSpec] = mysql_common.

            
                  ColumnSpecsBuilder.with_defaults().build(),

            
                  conflict_resolution: Optional[mysql_common.ConflictResolution] = None):

            
                self.connector_config = _MySQLConnectorConfig.from_base_config(

            
                    connection_config)

            
                super().__init__(

            
                    connection_config=self.connector_config.to_connection_config(),

            
                    write_config=write_config,

            
                    table_name=table_name,

            
                    column_specs=column_specs,

            
                    conflict_resolution=conflict_resolution)
	#
	# 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.

	from dataclasses import asdict
	from dataclasses import dataclass
	from typing import Any
	from typing import Optional

	from apache_beam.ml.rag.ingestion import mysql
	from apache_beam.ml.rag.ingestion import mysql_common
	from apache_beam.ml.rag.ingestion import postgres
	from apache_beam.ml.rag.ingestion import postgres_common
	from apache_beam.ml.rag.ingestion.jdbc_common import ConnectionConfig
	from apache_beam.ml.rag.ingestion.jdbc_common import WriteConfig


	@dataclass
	class LanguageConnectorConfig:
	"""Configuration options for CloudSQL Java language connector.

	Set parameters to connect connection to a CloudSQL instance using
	Java language connector connector. For details see
	https://github.com/GoogleCloudPlatform/cloud-sql-jdbc-socket-factory/blob/main/docs/jdbc.md

	Attributes:
	username: Database username.
	password: Database password. Can be empty string when using IAM.
	database_name: Name of the database to connect to.
	instance_name: Instance connection name. Format:
	'<PROJECT>:<REGION>:<INSTANCE>'
	ip_type: Preferred order of IP types used to connect via a comma
	list of strings.
	enable_iam_auth: Whether to enable IAM authentication. Default is False
	target_principal: Optional service account to impersonate for
	connection.
	delegates: Optional list of service accounts for delegated
	impersonation.
	admin_service_endpoint: Optional custom API service endpoint.
	quota_project: Optional project ID for quota and billing.
	connection_properties: Optional JDBC connection properties dict.
	Example: {'ssl': 'true'}
	additional_properties: Additional properties to be added to the JDBC
	url. Example: {'someProperty': 'true'}
	"""
	username: str
	password: str
	database_name: str
	instance_name: str
	ip_types: Optional[list[str]] = None
	enable_iam_auth: bool = False
	target_principal: Optional[str] = None
	delegates: Optional[list[str]] = None
	quota_project: Optional[str] = None
	connection_properties: Optional[dict[str, str]] = None
	additional_properties: Optional[dict[str, Any]] = None

	def _base_jdbc_properties(self) -> dict[str, Any]:
	properties = {"cloudSqlInstance": self.instance_name}

	if self.ip_types:
	properties["ipTypes"] = ",".join(self.ip_types)

	if self.enable_iam_auth:
	properties["enableIamAuth"] = "true"

	if self.target_principal:
	properties["cloudSqlTargetPrincipal"] = self.target_principal

	if self.delegates:
	properties["cloudSqlDelegates"] = ",".join(self.delegates)

	if self.quota_project:
	properties["cloudSqlAdminQuotaProject"] = self.quota_project

	if self.additional_properties:
	properties.update(self.additional_properties)

	return properties

	def _build_jdbc_url(self, socketFactory, database_type):
	url = f"jdbc:{database_type}:///{self.database_name}?"

	properties = self._base_jdbc_properties()
	properties['socketFactory'] = socketFactory

	property_string = "&".join(f"{k}={v}" for k, v in properties.items())
	return url + property_string

	def to_connection_config(self):
	return ConnectionConfig(
	jdbc_url=self.to_jdbc_url(),
	username=self.username,
	password=self.password,
	connection_properties=self.connection_properties,
	additional_jdbc_args=self.additional_jdbc_args())

	def additional_jdbc_args(self) -> dict[str, list[Any]]:
	return {}


	@dataclass
	class _PostgresConnectorConfig(LanguageConnectorConfig):
	def to_jdbc_url(self) -> str:
	"""Convert options to a properly formatted JDBC URL.

	Returns:
	JDBC URL string configured with all options.
	"""
	return self._build_jdbc_url(
	socketFactory="com.google.cloud.sql.postgres.SocketFactory",
	database_type="postgresql")

	def additional_jdbc_args(self) -> dict[str, list[Any]]:
	return {
	'classpath': [
	"org.postgresql:postgresql:42.2.16",
	"com.google.cloud.sql:postgres-socket-factory:1.25.0"
	]
	}

	@classmethod
	def from_base_config(cls, config: LanguageConnectorConfig):
	return cls(**asdict(config))


	class CloudSQLPostgresVectorWriterConfig(postgres.PostgresVectorWriterConfig):
	def __init__(
	self,
	connection_config: LanguageConnectorConfig,
	table_name: str,
	*,
	# pylint: disable=dangerous-default-value
	write_config: WriteConfig = WriteConfig(),
	column_specs: list[postgres_common.ColumnSpec] = postgres_common.
	ColumnSpecsBuilder.with_defaults().build(),
	conflict_resolution: Optional[
	postgres_common.ConflictResolution] = postgres_common.
	ConflictResolution(on_conflict_fields=[], action='IGNORE')):
	"""Configuration for writing vectors to ClouSQL Postgres.

	Supports flexible schema configuration through column specifications and
	conflict resolution strategies.

	Args:
	connection_config: :class:`LanguageConnectorConfig`.
	table_name: Target table name.
	write_config: JdbcIO :class:`~.jdbc_common.WriteConfig` to control
	batch sizes, authosharding, etc.
	column_specs:
	Use :class:`~.postgres_common.ColumnSpecsBuilder` to configure how
	embeddings and metadata are written a database
	schema. If None, uses default EmbeddableItem schema.
	conflict_resolution: Optional
	:class:`~.postgres_common.ConflictResolution`
	strategy for handling insert conflicts. ON CONFLICT DO NOTHING by
	default.

	Examples:
	Basic usage with default schema:

	>>> config = PostgresVectorWriterConfig(
	... connection_config=PostgresConnectionConfig(...),
	... table_name='embeddings'
	... )

	Simple case with default schema:

	>>> config = PostgresVectorWriterConfig(
	... connection_config=ConnectionConfig(...),
	... table_name='embeddings'
	... )

	Custom schema with metadata fields:

	>>> specs = (ColumnSpecsBuilder()
	... .with_id_spec(column_name="my_id_column")
	... .with_embedding_spec(column_name="embedding_vec")
	... .add_metadata_field(field="source", column_name="src")
	... .add_metadata_field(
	... "timestamp",
	... column_name="created_at",
	... sql_typecast="::timestamp"
	... )
	... .build())

	Minimal schema (only ID + embedding written)

	>>> column_specs = (ColumnSpecsBuilder()
	... .with_id_spec()
	... .with_embedding_spec()
	... .build())

	>>> config = CloudSQLPostgresVectorWriterConfig(
	... connection_config=PostgresConnectionConfig(...),
	... table_name='embeddings',
	... column_specs=specs
	... )
	"""
	self.connector_config = _PostgresConnectorConfig.from_base_config(
	connection_config)
	super().__init__(
	connection_config=self.connector_config.to_connection_config(),
	write_config=write_config,
	table_name=table_name,
	column_specs=column_specs,
	conflict_resolution=conflict_resolution)


	@dataclass
	class _MySQLConnectorConfig(LanguageConnectorConfig):
	def to_jdbc_url(self) -> str:
	"""Convert options to a properly formatted MySQL JDBC URL."""
	return self._build_jdbc_url(
	socketFactory="com.google.cloud.sql.mysql.SocketFactory",
	database_type="mysql")

	def additional_jdbc_args(self) -> dict[str, list[Any]]:
	return {
	'classpath': [
	"mysql:mysql-connector-java:8.0.22",
	"com.google.cloud.sql:mysql-socket-factory-connector-j-8:1.25.0"
	]
	}

	@classmethod
	def from_base_config(cls, config: LanguageConnectorConfig):
	return cls(**asdict(config))


	class CloudSQLMySQLVectorWriterConfig(mysql.MySQLVectorWriterConfig):
	def __init__(
	self,
	connection_config: LanguageConnectorConfig,
	table_name: str,
	*,
	write_config: WriteConfig = WriteConfig(),
	# pylint: disable=dangerous-default-value
	column_specs: list[mysql_common.ColumnSpec] = mysql_common.
	ColumnSpecsBuilder.with_defaults().build(),
	conflict_resolution: Optional[mysql_common.ConflictResolution] = None):
	self.connector_config = _MySQLConnectorConfig.from_base_config(
	connection_config)
	super().__init__(
	connection_config=self.connector_config.to_connection_config(),
	write_config=write_config,
	table_name=table_name,
	column_specs=column_specs,
	conflict_resolution=conflict_resolution)