spec

type AWSIAMAuthService string

type AWSIAMAuthSpec struct {
	// The AWS database service to authenticate to.
	Service AWSIAMAuthService `json:"service,omitempty" jsonschema:"enum=rds,default=rds"`

	// AWS region the database is in, for example `us-east-1`. If empty, the region
	// is resolved from the standard AWS configuration sources (the `AWS_REGION`
	// environment variable, the shared config file, ...).
	Region string `json:"region,omitempty" jsonschema:"example=us-east-1"`

	// The endpoint the token is signed for, as `host` or `host:port`. If empty, the
	// host and port from `connection_string` are used; if only a host is given, the
	// port from `connection_string` is used. Set this when the plugin connects
	// through a different address than the database endpoint itself, such as an SSH
	// tunnel or a CNAME.
	Endpoint string `json:"endpoint,omitempty" jsonschema:"example=mydb.123456789012.us-east-1.rds.amazonaws.com:5432"`

	// [Local profile](https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-files.html)
	// to use to authenticate with. This should be set to the name of the profile.
	LocalProfile string `json:"local_profile,omitempty" jsonschema:"example=my_aws_profile"`

	// If specified, this role is assumed and its credentials are used to sign the
	// authentication token.
	RoleARN string `json:"role_arn,omitempty" jsonschema:"pattern=^(arn(:[^:\n]*){5}([:/].*)?)?$"`

	// If specified, this session name is used when assuming `role_arn`.
	RoleSessionName string `json:"role_session_name,omitempty" jsonschema:"example=my_aws_role_session_name"`

	// If specified, this external ID is used when assuming `role_arn`.
	ExternalID string `json:"external_id,omitempty" jsonschema:"example=external_id"`
}

type LakebaseSpec struct {
	// The Lakebase database endpoint resource name, in the format
	// `projects/{project_id}/branches/{branch_id}/endpoints/{endpoint_id}`.
	Endpoint string `json:"endpoint" jsonschema:"required,minLength=1"`

	// Databricks workspace host, for example `https://your-workspace.cloud.databricks.com`.
	// If empty, the Databricks SDK resolves it from the `DATABRICKS_HOST` environment
	// variable (or other default Databricks configuration sources).
	Host string `json:"host,omitempty"`

	// Databricks service principal OAuth client ID.
	// If empty, the Databricks SDK resolves it from the `DATABRICKS_CLIENT_ID`
	// environment variable (or other default Databricks configuration sources).
	ClientID string `json:"client_id,omitempty"`

	// Databricks service principal OAuth client secret.
	// If empty, the Databricks SDK resolves it from the `DATABRICKS_CLIENT_SECRET`
	// environment variable (or other default Databricks configuration sources).
	ClientSecret string `json:"client_secret,omitempty"`
}

type LogLevel tracelog.LogLevel

type OpenAIEmbedding struct {
	APIKey     string `json:"api_key" jsonschema:"required,minLength=1,title=OpenAI API Key"`
	ModelName  string `json:"model_name" jsonschema:"required,minLength=1"`
	Dimensions int    `json:"dimensions" jsonschema:"minimum=1"`
}

type PgVectorConfig struct {
	// Tables to create embeddings for.
	Tables []PgVectorTableConfig `json:"tables,omitempty" jsonschema:"required,minItems=1"`
	// Optional text splitting configuration. If set, all sub-configurations must be set.
	TextSplitter *PgVectorTextSplitter `json:"text_splitter,omitempty"`
	// OpenAI embedding provider configuration. Required if PgVectorConfig is set.
	OpenAIEmbedding OpenAIEmbedding `json:"openai_embedding" jsonschema:"required"`
}

type PgVectorRecursiveText struct {
	ChunkSize    int `json:"chunk_size" jsonschema:"required,minimum=1"`
	ChunkOverlap int `json:"chunk_overlap" jsonschema:"required,minimum=0"`
}

type PgVectorTableConfig struct {
	SourceTableName string   `json:"source_table_name" jsonschema:"required,minLength=1"`
	TargetTableName string   `json:"target_table_name" jsonschema:"required,minLength=1"`
	EmbedColumns    []string `json:"embed_columns" jsonschema:"required,minItems=1"`
	MetadataColumns []string `json:"metadata_columns,omitempty"`
}

type PgVectorTextSplitter struct {
	RecursiveText PgVectorRecursiveText `json:"recursive_text" jsonschema:"required"`
}

type Spec struct {
	// Connection string to connect to the database. This can be a DSN or a URI, for example:
	//
	// - `"user=user password=pass host=localhost port=5432 dbname=mydb sslmode=disable"` DSN format
	// - `"postgres://user:pass@localhost:5432/mydb?sslmode=prefer"` connect with tcp and prefer TLS
	// - `"postgres://user:pass@localhost:5432/mydb?sslmode=disable&search_path=myschema"` connect with tcp, disable TLS and use a custom schema
	ConnectionString string `json:"connection_string,omitempty" jsonschema:"required,minLength=1,example=${POSTGRESQL_CONNECTION_STRING}"`

	// Available: `error`, `warn`, `info`, `debug`, `trace`.
	// Defines what [`pgx`](https://github.com/jackc/pgx) call events should be logged.
	PgxLogLevel LogLevel `json:"pgx_log_level,omitempty" jsonschema:"default=error"`

	// Maximum number of items that may be grouped together to be written in a single write.
	BatchSize int64 `json:"batch_size,omitempty" jsonschema:"minimum=1,default=10000"`

	// Maximum size of items that may be grouped together to be written in a single write.
	BatchSizeBytes int64 `json:"batch_size_bytes,omitempty" jsonschema:"minimum=1,default=100000000"`

	// Maximum interval between batch writes.
	BatchTimeout configtype.Duration `json:"batch_timeout,omitempty"`

	// Option to create specific indexes to improve deletion performance
	CreatePerformanceIndexes bool `json:"create_performance_indexes,omitempty" jsonschema:"default=false"`

	// Write rows with the PostgreSQL `COPY` protocol instead of `INSERT` statements.
	//
	// `COPY` sends a batch as a single statement rather than one statement per row,
	// which is substantially faster against a remote database. Tables with a
	// primary key are copied into a temporary staging table and merged into the
	// target, since `COPY` has no `ON CONFLICT` of its own.
	//
	// Rows repeating a primary key within one batch are collapsed before the merge,
	// keeping the last one, so a database-side trigger fires once rather than once
	// per row.
	//
	// Tables with a primary key are still written to by an `INSERT`, so row-level
	// security and rewrite rules behave as before. Tables without one are copied
	// straight in: rewrite rules are not applied, and row-level security makes the
	// write fail, since PostgreSQL rejects `COPY FROM` on such a table.
	//
	// It has no effect on CockroachDB, CrateDB, or when `pgvector_config` is set;
	// those always use `INSERT`.
	UseCopyFrom bool `json:"use_copy_from,omitempty" jsonschema:"default=false"`

	// Optional configuration to enable PgVector embedding support.
	PgVectorConfig *PgVectorConfig `json:"pgvector_config,omitempty"`

	// Number of times to retry a transaction if a deadlock is detected by Postgres.
	RetryOnDeadlock int64 `json:"retry_on_deadlock,omitempty" jsonschema:"default=0"`

	// Optional configuration to connect to [Databricks Lakebase](https://docs.databricks.com/aws/en/oltp),
	// a PostgreSQL-compatible managed database.
	Lakebase *LakebaseSpec `json:"lakebase,omitempty"`

	// Optional configuration to connect to an AWS-managed, PostgreSQL-compatible
	// database with [IAM database authentication](https://docs.aws.amazon.com/AmazonRDS/latest/UserGuide/UsingWithRDS.IAMDBAuth.html).
	AWSIAMAuth *AWSIAMAuthSpec `json:"aws_iam_auth,omitempty"`
}