From 1b7393647a18be2e2d2df7d3ad1ee96d502f2bb3 Mon Sep 17 00:00:00 2001
From: sanj <67624670+iodrift@users.noreply.github.com>
Date: Tue, 30 Jul 2024 16:34:25 -0700
Subject: [PATCH] Auto-update: Tue Jul 30 16:34:25 PDT 2024
---
sijapi/classes.py | 84 ++++++++++++++++++++++++++++++-----------------
1 file changed, 54 insertions(+), 30 deletions(-)
diff --git a/sijapi/classes.py b/sijapi/classes.py
index 546e78e..8eac7c6 100644
--- a/sijapi/classes.py
+++ b/sijapi/classes.py
@@ -437,37 +437,61 @@ class APIConfig(BaseModel):
)
""")
- await conn.execute(f"""
- DO $$
- BEGIN
- -- Ensure version column exists
- IF NOT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = '{table_name}' AND column_name = 'version') THEN
- ALTER TABLE "{table_name}" ADD COLUMN version INTEGER DEFAULT 1;
- END IF;
-
- -- Ensure server_id column exists
- IF NOT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = '{table_name}' AND column_name = 'server_id') THEN
- ALTER TABLE "{table_name}" ADD COLUMN server_id TEXT DEFAULT '{os.environ.get('TS_ID')}';
- END IF;
-
- -- Create or replace the trigger function
- CREATE OR REPLACE FUNCTION update_version_and_server_id()
- RETURNS TRIGGER AS $$
- BEGIN
- NEW.version = COALESCE(OLD.version, 0) + 1;
- NEW.server_id = '{os.environ.get('TS_ID')}';
- RETURN NEW;
- END;
- $$ LANGUAGE plpgsql;
-
- -- Create the trigger if it doesn't exist
- IF NOT EXISTS (SELECT 1 FROM pg_trigger WHERE tgname = 'update_version_and_server_id_trigger' AND tgrelid = '{table_name}'::regclass) THEN
- CREATE TRIGGER update_version_and_server_id_trigger
- BEFORE INSERT OR UPDATE ON "{table_name}"
- FOR EACH ROW EXECUTE FUNCTION update_version_and_server_id();
- END IF;
- END $$;
+ # Check if version column exists
+ version_exists = await conn.fetchval(f"""
+ SELECT EXISTS (
+ SELECT 1 FROM information_schema.columns
+ WHERE table_name = '{table_name}' AND column_name = 'version'
+ )
""")
+
+ # Check if server_id column exists
+ server_id_exists = await conn.fetchval(f"""
+ SELECT EXISTS (
+ SELECT 1 FROM information_schema.columns
+ WHERE table_name = '{table_name}' AND column_name = 'server_id'
+ )
+ """)
+
+ # Add version column if it doesn't exist
+ if not version_exists:
+ await conn.execute(f"""
+ ALTER TABLE "{table_name}" ADD COLUMN version INTEGER DEFAULT 1
+ """)
+
+ # Add server_id column if it doesn't exist
+ if not server_id_exists:
+ await conn.execute(f"""
+ ALTER TABLE "{table_name}" ADD COLUMN server_id TEXT DEFAULT '{os.environ.get('TS_ID')}'
+ """)
+
+ # Create or replace the trigger function
+ await conn.execute("""
+ CREATE OR REPLACE FUNCTION update_version_and_server_id()
+ RETURNS TRIGGER AS $$
+ BEGIN
+ NEW.version = COALESCE(OLD.version, 0) + 1;
+ NEW.server_id = $1;
+ RETURN NEW;
+ END;
+ $$ LANGUAGE plpgsql;
+ """)
+
+ # Create the trigger if it doesn't exist
+ trigger_exists = await conn.fetchval(f"""
+ SELECT EXISTS (
+ SELECT 1 FROM pg_trigger
+ WHERE tgname = 'update_version_and_server_id_trigger'
+ AND tgrelid = '{table_name}'::regclass
+ )
+ """)
+
+ if not trigger_exists:
+ await conn.execute(f"""
+ CREATE TRIGGER update_version_and_server_id_trigger
+ BEFORE INSERT OR UPDATE ON "{table_name}"
+ FOR EACH ROW EXECUTE FUNCTION update_version_and_server_id('{os.environ.get('TS_ID')}')
+ """)
info(f"Successfully ensured sync columns and trigger for table {table_name}. Has primary key: {has_primary_key}")
return has_primary_key