fix: DB migration conflict (#3739)
This commit is contained in:
+60
-20
@@ -100,8 +100,8 @@ def insert_to_pg(query, data):
|
|||||||
logger.error(exc)
|
logger.error(exc)
|
||||||
logger.error(f"Failed to insert {d}")
|
logger.error(f"Failed to insert {d}")
|
||||||
else:
|
else:
|
||||||
logger.error("query:", query)
|
logger.error("query: " + query)
|
||||||
logger.error("data:", d)
|
logger.error("data: " + str(d))
|
||||||
raise ValueError(f"Failed to insert {d}") from exc
|
raise ValueError(f"Failed to insert {d}") from exc
|
||||||
connection.commit()
|
connection.commit()
|
||||||
|
|
||||||
@@ -125,6 +125,7 @@ def migrate_ext(file: str):
|
|||||||
migrate_db(file, schema)
|
migrate_db(file, schema)
|
||||||
logger.info(f"✅ Migrated ext: {schema}")
|
logger.info(f"✅ Migrated ext: {schema}")
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
|
logger.error(exc)
|
||||||
logger.error(f"🛑 Failed to migrate extension {schema}: {exc}")
|
logger.error(f"🛑 Failed to migrate extension {schema}: {exc}")
|
||||||
|
|
||||||
|
|
||||||
@@ -134,8 +135,8 @@ def migrate_db(file: str, schema: str, exclude_tables: list[str] | None = None):
|
|||||||
exclude_tables = []
|
exclude_tables = []
|
||||||
assert os.path.isfile(file), f"{file} does not exist!"
|
assert os.path.isfile(file), f"{file} does not exist!"
|
||||||
|
|
||||||
cursor = get_sqlite_cursor(file)
|
sqlite_cursor = get_sqlite_cursor(file)
|
||||||
tables = cursor.execute(
|
tables = sqlite_cursor.execute(
|
||||||
"""
|
"""
|
||||||
SELECT name FROM sqlite_master
|
SELECT name FROM sqlite_master
|
||||||
WHERE type='table' AND name not like 'sqlite?_%' escape '?'
|
WHERE type='table' AND name not like 'sqlite?_%' escape '?'
|
||||||
@@ -151,16 +152,18 @@ def migrate_db(file: str, schema: str, exclude_tables: list[str] | None = None):
|
|||||||
if exclude_tables and table_name in exclude_tables:
|
if exclude_tables and table_name in exclude_tables:
|
||||||
continue
|
continue
|
||||||
|
|
||||||
columns = cursor.execute(f"PRAGMA table_info({table_name})").fetchall()
|
columns = build_table_columns(file, schema, table_name)
|
||||||
q = build_insert_query(schema, table_name, columns)
|
q = build_insert_query(schema, table_name, columns)
|
||||||
|
|
||||||
data = cursor.execute(f"SELECT * FROM {table_name};").fetchall()
|
data = sqlite_cursor.execute(f"SELECT * FROM {table_name};").fetchall()
|
||||||
|
|
||||||
if len(data) == 0:
|
if len(data) == 0:
|
||||||
logger.warning(f"🛑 You sneaky dev! Table {table_name} is empty!")
|
logger.warning(f"⚠️ You sneaky dev! Table {table_name} is empty!")
|
||||||
|
continue
|
||||||
|
|
||||||
insert_to_pg(q, data)
|
insert_to_pg(q, data)
|
||||||
cursor.close()
|
logger.info(f"✅ Migrated table '{schema}.{table_name}' successfully")
|
||||||
|
sqlite_cursor.close()
|
||||||
|
|
||||||
|
|
||||||
def build_insert_query(schema, table_name, columns):
|
def build_insert_query(schema, table_name, columns):
|
||||||
@@ -174,6 +177,30 @@ def build_insert_query(schema, table_name, columns):
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
def build_table_columns(file: str, schema: str, table_name: str):
|
||||||
|
sqlite_cursor = get_sqlite_cursor(file)
|
||||||
|
pg_cursor = get_postgres_cursor()
|
||||||
|
|
||||||
|
sqlite_columns = sqlite_cursor.execute(
|
||||||
|
f"PRAGMA table_info({table_name})"
|
||||||
|
).fetchall()
|
||||||
|
pg_cursor.execute(
|
||||||
|
f"""
|
||||||
|
SELECT table_name, column_name, udt_name FROM information_schema.columns
|
||||||
|
WHERE table_schema = '{schema}'AND table_name = '{table_name}';"""
|
||||||
|
)
|
||||||
|
pg_columns = pg_cursor.fetchall()
|
||||||
|
|
||||||
|
columns = []
|
||||||
|
for sqlite_col in sqlite_columns:
|
||||||
|
for pg_col in pg_columns:
|
||||||
|
if sqlite_col[1].lower() == pg_col[1].lower():
|
||||||
|
columns.append((sqlite_col[0], sqlite_col[1], pg_col[2]))
|
||||||
|
break
|
||||||
|
sqlite_cursor.close()
|
||||||
|
return columns
|
||||||
|
|
||||||
|
|
||||||
def build_on_conflict_query_statement(schema, table_name, columns):
|
def build_on_conflict_query_statement(schema, table_name, columns):
|
||||||
unique_cols = table_unique_columns(schema, table_name)
|
unique_cols = table_unique_columns(schema, table_name)
|
||||||
if len(unique_cols) == 0:
|
if len(unique_cols) == 0:
|
||||||
@@ -191,23 +218,36 @@ def build_on_conflict_query_statement(schema, table_name, columns):
|
|||||||
def table_unique_columns(schema, table_name):
|
def table_unique_columns(schema, table_name):
|
||||||
cursor = get_postgres_cursor()
|
cursor = get_postgres_cursor()
|
||||||
query = f"""
|
query = f"""
|
||||||
SELECT a.attname
|
SELECT
|
||||||
FROM pg_index i
|
array_agg(a.attname ORDER BY a.attnum) AS columns,
|
||||||
JOIN pg_attribute a ON a.attrelid = i.indrelid
|
i.indisprimary as is_primary,
|
||||||
AND a.attnum = ANY(i.indkey)
|
i.indexrelid::regclass AS index_name,
|
||||||
WHERE i.indrelid = '{schema}.{table_name}'::regclass
|
COUNT(*) AS column_count,
|
||||||
AND i.indisunique;
|
(COUNT(*) = 1) AS is_individual
|
||||||
|
FROM pg_index i
|
||||||
|
JOIN pg_attribute a
|
||||||
|
ON a.attrelid = i.indrelid
|
||||||
|
AND a.attnum = ANY (i.indkey)
|
||||||
|
WHERE i.indrelid = '{schema}.{table_name}'::regclass
|
||||||
|
AND i.indisunique
|
||||||
|
GROUP BY i.indexrelid;
|
||||||
"""
|
"""
|
||||||
cursor.execute(query)
|
cursor.execute(query)
|
||||||
columns = [row[0] for row in cursor.fetchall()]
|
rows = cursor.fetchall()
|
||||||
|
columns = [row[0] for row in rows if not row[1]] # exclude primary keys
|
||||||
|
if len(columns) == 0:
|
||||||
|
# use primary keys if no unique keys found
|
||||||
|
columns = [row[0] for row in rows if row[1]]
|
||||||
cursor.close()
|
cursor.close()
|
||||||
return columns
|
if len(columns) == 0:
|
||||||
|
return []
|
||||||
|
return columns[0]
|
||||||
|
|
||||||
|
|
||||||
def to_column_type(column_type):
|
def to_column_type(column_type: str):
|
||||||
if column_type == "TIMESTAMP":
|
if column_type.upper() == "TIMESTAMP":
|
||||||
return "to_timestamp(%s)"
|
return "to_timestamp(%s)"
|
||||||
if column_type in ["BOOLEAN", "BOOL"]:
|
if column_type.upper() in ["BOOLEAN", "BOOL"]:
|
||||||
return "%s::boolean"
|
return "%s::boolean"
|
||||||
return "%s"
|
return "%s"
|
||||||
|
|
||||||
@@ -255,7 +295,7 @@ parser.add_argument(
|
|||||||
|
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
|
|
||||||
logger.info("Selected path: ", args.sqlite_path)
|
logger.info("Selected path: " + args.sqlite_path)
|
||||||
|
|
||||||
if os.path.isdir(args.sqlite_path):
|
if os.path.isdir(args.sqlite_path):
|
||||||
exclude_tables = ["dbversions"]
|
exclude_tables = ["dbversions"]
|
||||||
|
|||||||
Reference in New Issue
Block a user