Translate commnet to English
Browse files- lightrag/kg/tidb_impl.py +16 -16
lightrag/kg/tidb_impl.py
CHANGED
@@ -45,8 +45,8 @@ class TiDB:
|
|
45 |
raise
|
46 |
|
47 |
async def _migrate_timestamp_columns(self):
|
48 |
-
"""
|
49 |
-
#
|
50 |
tables_to_migrate = {
|
51 |
"LIGHTRAG_GRAPH_NODES": ["createtime", "updatetime"],
|
52 |
"LIGHTRAG_GRAPH_EDGES": ["createtime", "updatetime"],
|
@@ -56,7 +56,7 @@ class TiDB:
|
|
56 |
for table_name, columns in tables_to_migrate.items():
|
57 |
for column_name in columns:
|
58 |
try:
|
59 |
-
#
|
60 |
check_column_sql = f"""
|
61 |
SELECT COLUMN_NAME, DATA_TYPE, COLUMN_TYPE
|
62 |
FROM INFORMATION_SCHEMA.COLUMNS
|
@@ -67,27 +67,27 @@ class TiDB:
|
|
67 |
column_info = await self.query(check_column_sql)
|
68 |
if not column_info:
|
69 |
logger.warning(
|
70 |
-
f"
|
71 |
)
|
72 |
continue
|
73 |
|
74 |
-
#
|
75 |
data_type = column_info.get("DATA_TYPE", "").lower()
|
76 |
column_type = column_info.get("COLUMN_TYPE", "").lower()
|
77 |
|
78 |
-
#
|
79 |
if data_type == "timestamp" and "time zone" in column_type:
|
80 |
logger.info(
|
81 |
-
f"
|
82 |
)
|
83 |
continue
|
84 |
|
85 |
-
#
|
86 |
if data_type == "datetime" or (
|
87 |
data_type == "timestamp" and "time zone" not in column_type
|
88 |
):
|
89 |
logger.info(
|
90 |
-
f"
|
91 |
)
|
92 |
migration_sql = f"""
|
93 |
ALTER TABLE {table_name}
|
@@ -96,14 +96,14 @@ class TiDB:
|
|
96 |
|
97 |
await self.execute(migration_sql)
|
98 |
logger.info(
|
99 |
-
f"
|
100 |
)
|
101 |
except Exception as e:
|
102 |
-
#
|
103 |
-
logger.warning(f"
|
104 |
|
105 |
async def check_tables(self):
|
106 |
-
#
|
107 |
for k, v in TABLES.items():
|
108 |
try:
|
109 |
await self.query(f"SELECT 1 FROM {k}".format(k=k))
|
@@ -117,12 +117,12 @@ class TiDB:
|
|
117 |
logger.error(f"Failed to create table {k} in TiDB database")
|
118 |
logger.error(f"TiDB database error: {e}")
|
119 |
|
120 |
-
#
|
121 |
try:
|
122 |
await self._migrate_timestamp_columns()
|
123 |
except Exception as e:
|
124 |
logger.error(f"TiDB, Failed to migrate timestamp columns: {e}")
|
125 |
-
#
|
126 |
|
127 |
async def query(
|
128 |
self, sql: str, params: dict = None, multirows: bool = False
|
@@ -397,7 +397,7 @@ class TiDBKVStorage(BaseKVStorage):
|
|
397 |
if table_name != "LIGHTRAG_LLM_CACHE":
|
398 |
return False
|
399 |
|
400 |
-
#
|
401 |
modes_list = ", ".join([f"'{mode}'" for mode in modes])
|
402 |
sql = f"""
|
403 |
DELETE FROM {table_name}
|
|
|
45 |
raise
|
46 |
|
47 |
async def _migrate_timestamp_columns(self):
|
48 |
+
"""Migrate timestamp columns in tables to timezone-aware types, assuming original data is in UTC"""
|
49 |
+
# Tables and columns that need migration
|
50 |
tables_to_migrate = {
|
51 |
"LIGHTRAG_GRAPH_NODES": ["createtime", "updatetime"],
|
52 |
"LIGHTRAG_GRAPH_EDGES": ["createtime", "updatetime"],
|
|
|
56 |
for table_name, columns in tables_to_migrate.items():
|
57 |
for column_name in columns:
|
58 |
try:
|
59 |
+
# Check if column exists
|
60 |
check_column_sql = f"""
|
61 |
SELECT COLUMN_NAME, DATA_TYPE, COLUMN_TYPE
|
62 |
FROM INFORMATION_SCHEMA.COLUMNS
|
|
|
67 |
column_info = await self.query(check_column_sql)
|
68 |
if not column_info:
|
69 |
logger.warning(
|
70 |
+
f"Column {table_name}.{column_name} does not exist, skipping migration"
|
71 |
)
|
72 |
continue
|
73 |
|
74 |
+
# Check column type
|
75 |
data_type = column_info.get("DATA_TYPE", "").lower()
|
76 |
column_type = column_info.get("COLUMN_TYPE", "").lower()
|
77 |
|
78 |
+
# If already timestamp type, check if it contains timezone information
|
79 |
if data_type == "timestamp" and "time zone" in column_type:
|
80 |
logger.info(
|
81 |
+
f"Column {table_name}.{column_name} is already a timezone-aware timestamp type, no migration needed"
|
82 |
)
|
83 |
continue
|
84 |
|
85 |
+
# If datetime type, need to migrate to timestamp
|
86 |
if data_type == "datetime" or (
|
87 |
data_type == "timestamp" and "time zone" not in column_type
|
88 |
):
|
89 |
logger.info(
|
90 |
+
f"Migrating {table_name}.{column_name} to timestamp type"
|
91 |
)
|
92 |
migration_sql = f"""
|
93 |
ALTER TABLE {table_name}
|
|
|
96 |
|
97 |
await self.execute(migration_sql)
|
98 |
logger.info(
|
99 |
+
f"Successfully migrated {table_name}.{column_name} to timestamp type"
|
100 |
)
|
101 |
except Exception as e:
|
102 |
+
# Log error but don't interrupt the process
|
103 |
+
logger.warning(f"Failed to migrate {table_name}.{column_name}: {e}")
|
104 |
|
105 |
async def check_tables(self):
|
106 |
+
# First create all tables
|
107 |
for k, v in TABLES.items():
|
108 |
try:
|
109 |
await self.query(f"SELECT 1 FROM {k}".format(k=k))
|
|
|
117 |
logger.error(f"Failed to create table {k} in TiDB database")
|
118 |
logger.error(f"TiDB database error: {e}")
|
119 |
|
120 |
+
# After all tables are created, try to migrate timestamp fields
|
121 |
try:
|
122 |
await self._migrate_timestamp_columns()
|
123 |
except Exception as e:
|
124 |
logger.error(f"TiDB, Failed to migrate timestamp columns: {e}")
|
125 |
+
# Don't raise exceptions, allow initialization process to continue
|
126 |
|
127 |
async def query(
|
128 |
self, sql: str, params: dict = None, multirows: bool = False
|
|
|
397 |
if table_name != "LIGHTRAG_LLM_CACHE":
|
398 |
return False
|
399 |
|
400 |
+
# Build MySQL style IN query
|
401 |
modes_list = ", ".join([f"'{mode}'" for mode in modes])
|
402 |
sql = f"""
|
403 |
DELETE FROM {table_name}
|