replace db with sa to get typing support (#23240)
This commit is contained in:
@@ -2,11 +2,11 @@ from datetime import datetime
|
||||
from enum import Enum
|
||||
from typing import Optional
|
||||
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy import DateTime, String, func, text
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from .base import Base
|
||||
from .engine import db
|
||||
from .types import StringUUID
|
||||
|
||||
|
||||
@@ -47,9 +47,9 @@ class Provider(Base):
|
||||
|
||||
__tablename__ = "providers"
|
||||
__table_args__ = (
|
||||
db.PrimaryKeyConstraint("id", name="provider_pkey"),
|
||||
db.Index("provider_tenant_id_provider_idx", "tenant_id", "provider_name"),
|
||||
db.UniqueConstraint(
|
||||
sa.PrimaryKeyConstraint("id", name="provider_pkey"),
|
||||
sa.Index("provider_tenant_id_provider_idx", "tenant_id", "provider_name"),
|
||||
sa.UniqueConstraint(
|
||||
"tenant_id", "provider_name", "provider_type", "quota_type", name="unique_provider_name_type_quota"
|
||||
),
|
||||
)
|
||||
@@ -60,15 +60,15 @@ class Provider(Base):
|
||||
provider_type: Mapped[str] = mapped_column(
|
||||
String(40), nullable=False, server_default=text("'custom'::character varying")
|
||||
)
|
||||
encrypted_config: Mapped[Optional[str]] = mapped_column(db.Text, nullable=True)
|
||||
is_valid: Mapped[bool] = mapped_column(db.Boolean, nullable=False, server_default=text("false"))
|
||||
encrypted_config: Mapped[Optional[str]] = mapped_column(sa.Text, nullable=True)
|
||||
is_valid: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=text("false"))
|
||||
last_used: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
|
||||
|
||||
quota_type: Mapped[Optional[str]] = mapped_column(
|
||||
String(40), nullable=True, server_default=text("''::character varying")
|
||||
)
|
||||
quota_limit: Mapped[Optional[int]] = mapped_column(db.BigInteger, nullable=True)
|
||||
quota_used: Mapped[Optional[int]] = mapped_column(db.BigInteger, default=0)
|
||||
quota_limit: Mapped[Optional[int]] = mapped_column(sa.BigInteger, nullable=True)
|
||||
quota_used: Mapped[Optional[int]] = mapped_column(sa.BigInteger, default=0)
|
||||
|
||||
created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
|
||||
updated_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
|
||||
@@ -104,9 +104,9 @@ class ProviderModel(Base):
|
||||
|
||||
__tablename__ = "provider_models"
|
||||
__table_args__ = (
|
||||
db.PrimaryKeyConstraint("id", name="provider_model_pkey"),
|
||||
db.Index("provider_model_tenant_id_provider_idx", "tenant_id", "provider_name"),
|
||||
db.UniqueConstraint(
|
||||
sa.PrimaryKeyConstraint("id", name="provider_model_pkey"),
|
||||
sa.Index("provider_model_tenant_id_provider_idx", "tenant_id", "provider_name"),
|
||||
sa.UniqueConstraint(
|
||||
"tenant_id", "provider_name", "model_name", "model_type", name="unique_provider_model_name"
|
||||
),
|
||||
)
|
||||
@@ -116,8 +116,8 @@ class ProviderModel(Base):
|
||||
provider_name: Mapped[str] = mapped_column(String(255), nullable=False)
|
||||
model_name: Mapped[str] = mapped_column(String(255), nullable=False)
|
||||
model_type: Mapped[str] = mapped_column(String(40), nullable=False)
|
||||
encrypted_config: Mapped[Optional[str]] = mapped_column(db.Text, nullable=True)
|
||||
is_valid: Mapped[bool] = mapped_column(db.Boolean, nullable=False, server_default=text("false"))
|
||||
encrypted_config: Mapped[Optional[str]] = mapped_column(sa.Text, nullable=True)
|
||||
is_valid: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=text("false"))
|
||||
created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
|
||||
updated_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
|
||||
|
||||
@@ -125,8 +125,8 @@ class ProviderModel(Base):
|
||||
class TenantDefaultModel(Base):
|
||||
__tablename__ = "tenant_default_models"
|
||||
__table_args__ = (
|
||||
db.PrimaryKeyConstraint("id", name="tenant_default_model_pkey"),
|
||||
db.Index("tenant_default_model_tenant_id_provider_type_idx", "tenant_id", "provider_name", "model_type"),
|
||||
sa.PrimaryKeyConstraint("id", name="tenant_default_model_pkey"),
|
||||
sa.Index("tenant_default_model_tenant_id_provider_type_idx", "tenant_id", "provider_name", "model_type"),
|
||||
)
|
||||
|
||||
id: Mapped[str] = mapped_column(StringUUID, server_default=text("uuid_generate_v4()"))
|
||||
@@ -141,8 +141,8 @@ class TenantDefaultModel(Base):
|
||||
class TenantPreferredModelProvider(Base):
|
||||
__tablename__ = "tenant_preferred_model_providers"
|
||||
__table_args__ = (
|
||||
db.PrimaryKeyConstraint("id", name="tenant_preferred_model_provider_pkey"),
|
||||
db.Index("tenant_preferred_model_provider_tenant_provider_idx", "tenant_id", "provider_name"),
|
||||
sa.PrimaryKeyConstraint("id", name="tenant_preferred_model_provider_pkey"),
|
||||
sa.Index("tenant_preferred_model_provider_tenant_provider_idx", "tenant_id", "provider_name"),
|
||||
)
|
||||
|
||||
id: Mapped[str] = mapped_column(StringUUID, server_default=text("uuid_generate_v4()"))
|
||||
@@ -156,8 +156,8 @@ class TenantPreferredModelProvider(Base):
|
||||
class ProviderOrder(Base):
|
||||
__tablename__ = "provider_orders"
|
||||
__table_args__ = (
|
||||
db.PrimaryKeyConstraint("id", name="provider_order_pkey"),
|
||||
db.Index("provider_order_tenant_provider_idx", "tenant_id", "provider_name"),
|
||||
sa.PrimaryKeyConstraint("id", name="provider_order_pkey"),
|
||||
sa.Index("provider_order_tenant_provider_idx", "tenant_id", "provider_name"),
|
||||
)
|
||||
|
||||
id: Mapped[str] = mapped_column(StringUUID, server_default=text("uuid_generate_v4()"))
|
||||
@@ -167,9 +167,9 @@ class ProviderOrder(Base):
|
||||
payment_product_id: Mapped[str] = mapped_column(String(191), nullable=False)
|
||||
payment_id: Mapped[Optional[str]] = mapped_column(String(191))
|
||||
transaction_id: Mapped[Optional[str]] = mapped_column(String(191))
|
||||
quantity: Mapped[int] = mapped_column(db.Integer, nullable=False, server_default=text("1"))
|
||||
quantity: Mapped[int] = mapped_column(sa.Integer, nullable=False, server_default=text("1"))
|
||||
currency: Mapped[Optional[str]] = mapped_column(String(40))
|
||||
total_amount: Mapped[Optional[int]] = mapped_column(db.Integer)
|
||||
total_amount: Mapped[Optional[int]] = mapped_column(sa.Integer)
|
||||
payment_status: Mapped[str] = mapped_column(
|
||||
String(40), nullable=False, server_default=text("'wait_pay'::character varying")
|
||||
)
|
||||
@@ -187,8 +187,8 @@ class ProviderModelSetting(Base):
|
||||
|
||||
__tablename__ = "provider_model_settings"
|
||||
__table_args__ = (
|
||||
db.PrimaryKeyConstraint("id", name="provider_model_setting_pkey"),
|
||||
db.Index("provider_model_setting_tenant_provider_model_idx", "tenant_id", "provider_name", "model_type"),
|
||||
sa.PrimaryKeyConstraint("id", name="provider_model_setting_pkey"),
|
||||
sa.Index("provider_model_setting_tenant_provider_model_idx", "tenant_id", "provider_name", "model_type"),
|
||||
)
|
||||
|
||||
id: Mapped[str] = mapped_column(StringUUID, server_default=text("uuid_generate_v4()"))
|
||||
@@ -196,8 +196,8 @@ class ProviderModelSetting(Base):
|
||||
provider_name: Mapped[str] = mapped_column(String(255), nullable=False)
|
||||
model_name: Mapped[str] = mapped_column(String(255), nullable=False)
|
||||
model_type: Mapped[str] = mapped_column(String(40), nullable=False)
|
||||
enabled: Mapped[bool] = mapped_column(db.Boolean, nullable=False, server_default=text("true"))
|
||||
load_balancing_enabled: Mapped[bool] = mapped_column(db.Boolean, nullable=False, server_default=text("false"))
|
||||
enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=text("true"))
|
||||
load_balancing_enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=text("false"))
|
||||
created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
|
||||
updated_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
|
||||
|
||||
@@ -209,8 +209,8 @@ class LoadBalancingModelConfig(Base):
|
||||
|
||||
__tablename__ = "load_balancing_model_configs"
|
||||
__table_args__ = (
|
||||
db.PrimaryKeyConstraint("id", name="load_balancing_model_config_pkey"),
|
||||
db.Index("load_balancing_model_config_tenant_provider_model_idx", "tenant_id", "provider_name", "model_type"),
|
||||
sa.PrimaryKeyConstraint("id", name="load_balancing_model_config_pkey"),
|
||||
sa.Index("load_balancing_model_config_tenant_provider_model_idx", "tenant_id", "provider_name", "model_type"),
|
||||
)
|
||||
|
||||
id: Mapped[str] = mapped_column(StringUUID, server_default=text("uuid_generate_v4()"))
|
||||
@@ -219,7 +219,7 @@ class LoadBalancingModelConfig(Base):
|
||||
model_name: Mapped[str] = mapped_column(String(255), nullable=False)
|
||||
model_type: Mapped[str] = mapped_column(String(40), nullable=False)
|
||||
name: Mapped[str] = mapped_column(String(255), nullable=False)
|
||||
encrypted_config: Mapped[Optional[str]] = mapped_column(db.Text, nullable=True)
|
||||
enabled: Mapped[bool] = mapped_column(db.Boolean, nullable=False, server_default=text("true"))
|
||||
encrypted_config: Mapped[Optional[str]] = mapped_column(sa.Text, nullable=True)
|
||||
enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=text("true"))
|
||||
created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
|
||||
updated_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
|
||||
|
Reference in New Issue
Block a user