From 798f692e0bf74e3fe5c1526bb71bbd90f0ffc0d9 Mon Sep 17 00:00:00 2001
From: zxstty <zhaojiaqi18@huawei.com>
Date: Sun, 27 Apr 2025 15:43:05 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0opengauss=E5=90=91=E9=87=8F?=
 =?UTF-8?q?=E5=8C=96=E7=B4=A2=E5=BC=95=E7=9A=84=E9=80=82=E9=85=8D?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 chat2db/database/postgres.py | 59 +++++++++++++++++++++++++-----------
 1 file changed, 41 insertions(+), 18 deletions(-)

diff --git a/chat2db/database/postgres.py b/chat2db/database/postgres.py
index c3b8bf84..380eb3a3 100644
--- a/chat2db/database/postgres.py
+++ b/chat2db/database/postgres.py
@@ -33,15 +33,26 @@ class TableInfo(Base):
         TIMESTAMP(timezone=True),
         server_default=func.current_timestamp(),
         onupdate=func.current_timestamp())
-    __table_args__ = (
-        Index(
-            'table_note_vector_index',
-            table_note_vector,
-            postgresql_using='hnsw',
-            postgresql_with={'m': 16, 'ef_construction': 200},
-            postgresql_ops={'table_note_vector': 'vector_cosine_ops'}
-        ),
-    )
+    if 'opengauss' in config['DATABASE_URL']:
+        __table_args__ = (
+            Index(
+                'table_note_vector_index',
+                table_note_vector,
+                opengauss_using='hnsw',
+                opengauss_with={'m': 16, 'ef_construction': 200},
+                opengauss_ops={'vector': 'vector_cosine_ops'}
+            ),
+        )
+    else:
+        __table_args__ = (
+            Index(
+                'table_note_vector_index',
+                table_note_vector,
+                postgresql_using='hnsw',
+                postgresql_with={'m': 16, 'ef_construction': 200},
+                postgresql_ops={'table_note_vector': 'vector_cosine_ops'}
+            ),
+        )
 
 
 class ColumnInfo(Base):
@@ -66,15 +77,26 @@ class SqlExample(Base):
         TIMESTAMP(timezone=True),
         server_default=func.current_timestamp(),
         onupdate=func.current_timestamp())
-    __table_args__ = (
-        Index(
-            'question_vector_index',
-            question_vector,
-            postgresql_using='hnsw',
-            postgresql_with={'m': 16, 'ef_construction': 200},
-            postgresql_ops={'question_vector': 'vector_cosine_ops'}
-        ),
-    )
+    if 'opengauss' in config['DATABASE_URL']:
+        __table_args__ = (
+            Index(
+                'question_vector_index',
+                question_vector,
+                opengauss_using='hnsw',
+                opengauss_with={'m': 16, 'ef_construction': 200},
+                opengauss_ops={'vector': 'vector_cosine_ops'}
+            ),
+        )
+    else:
+        __table_args__ = (
+            Index(
+                'question_vector_index',
+                question_vector,
+                postgresql_using='hnsw',
+                postgresql_with={'m': 16, 'ef_construction': 200},
+                postgresql_ops={'question_vector': 'vector_cosine_ops'}
+            ),
+        )
 
 
 class PostgresDB:
@@ -94,6 +116,7 @@ class PostgresDB:
             if 'opengauss' in config['DATABASE_URL']:
                 from sqlalchemy import event
                 from opengauss_sqlalchemy.register_async import register_vector
+
                 @event.listens_for(cls.engine.sync_engine, "connect")
                 def connect(dbapi_connection, connection_record):
                     dbapi_connection.run_async(register_vector)
-- 
Gitee