Much more efficient database sync method, initial attempt.
This commit is contained in:
sanj
parent
011673893d
commit
d01f24ad45
3 changed files with 199 additions and 303 deletions
|
|
@ -134,6 +134,17 @@ async def handle_exception_middleware(request: Request, call_next):
|
|||
raise
|
||||
return response
|
||||
|
||||
|
||||
# This was removed on 7/31/2024 when we decided to instead use a targeted push sync approach.
|
||||
deprecated = '''
|
||||
async def push_changes_background():
|
||||
try:
|
||||
await API.push_changes_to_all()
|
||||
except Exception as e:
|
||||
err(f"Error pushing changes to other databases: {str(e)}")
|
||||
err(f"Traceback: {traceback.format_exc()}")
|
||||
|
||||
|
||||
@app.middleware("http")
|
||||
async def sync_middleware(request: Request, call_next):
|
||||
response = await call_next(request)
|
||||
|
|
@ -147,6 +158,7 @@ async def sync_middleware(request: Request, call_next):
|
|||
err(f"Error pushing changes to other databases: {str(e)}")
|
||||
|
||||
return response
|
||||
'''
|
||||
|
||||
def load_router(router_name):
|
||||
router_file = ROUTER_DIR / f'{router_name}.py'
|
||||
|
|
|
|||
|
|
@ -167,7 +167,6 @@ class Configuration(BaseModel):
|
|||
|
||||
|
||||
|
||||
|
||||
class APIConfig(BaseModel):
|
||||
HOST: str
|
||||
PORT: int
|
||||
|
|
@ -338,55 +337,6 @@ class APIConfig(BaseModel):
|
|||
err(f"Failed to acquire connection from pool for {pool_key}: {str(e)}")
|
||||
yield None
|
||||
|
||||
async def push_changes_to_one(self, pool_entry):
|
||||
try:
|
||||
async with self.get_connection() as local_conn:
|
||||
if local_conn is None:
|
||||
err(f"Failed to connect to local database. Skipping push to {pool_entry['ts_id']}")
|
||||
return
|
||||
|
||||
async with self.get_connection(pool_entry) as remote_conn:
|
||||
if remote_conn is None:
|
||||
err(f"Failed to connect to remote database {pool_entry['ts_id']}. Skipping push.")
|
||||
return
|
||||
|
||||
tables = await local_conn.fetch("""
|
||||
SELECT tablename FROM pg_tables
|
||||
WHERE schemaname = 'public'
|
||||
""")
|
||||
|
||||
for table in tables:
|
||||
table_name = table['tablename']
|
||||
try:
|
||||
if table_name in self.SPECIAL_TABLES:
|
||||
await self.sync_special_table(local_conn, remote_conn, table_name)
|
||||
else:
|
||||
primary_key = await self.ensure_sync_columns(remote_conn, table_name)
|
||||
last_synced_version = await self.get_last_synced_version(remote_conn, table_name, os.environ.get('TS_ID'))
|
||||
|
||||
changes = await local_conn.fetch(f"""
|
||||
SELECT * FROM "{table_name}"
|
||||
WHERE version > $1 AND server_id = $2
|
||||
ORDER BY version ASC
|
||||
""", last_synced_version, os.environ.get('TS_ID'))
|
||||
|
||||
if changes:
|
||||
changes_count = await self.apply_batch_changes(remote_conn, table_name, changes, primary_key)
|
||||
|
||||
if changes_count > 0:
|
||||
debug(f"Pushed {changes_count} changes for table {table_name} to {pool_entry['ts_id']}")
|
||||
|
||||
except Exception as e:
|
||||
err(f"Error pushing changes for table {table_name} to {pool_entry['ts_id']}: {str(e)}")
|
||||
err(f"Traceback: {traceback.format_exc()}")
|
||||
|
||||
info(f"Successfully pushed changes to {pool_entry['ts_id']}")
|
||||
|
||||
except Exception as e:
|
||||
err(f"Error pushing changes to {pool_entry['ts_id']}: {str(e)}")
|
||||
err(f"Traceback: {traceback.format_exc()}")
|
||||
|
||||
|
||||
async def close_db_pools(self):
|
||||
info("Closing database connection pools...")
|
||||
for pool_key, pool in self.db_pools.items():
|
||||
|
|
@ -500,114 +450,6 @@ class APIConfig(BaseModel):
|
|||
err(f"Error ensuring sync columns for table {table_name}: {str(e)}")
|
||||
err(f"Traceback: {traceback.format_exc()}")
|
||||
|
||||
async def apply_batch_changes(self, conn, table_name, changes, primary_key):
|
||||
if conn is None or not changes:
|
||||
debug(f"Skipping apply_batch_changes because conn is none or there are no changes.")
|
||||
return 0
|
||||
|
||||
try:
|
||||
columns = list(changes[0].keys())
|
||||
placeholders = [f'${i+1}' for i in range(len(columns))]
|
||||
|
||||
if primary_key:
|
||||
insert_query = f"""
|
||||
INSERT INTO "{table_name}" ({', '.join(f'"{col}"' for col in columns)})
|
||||
VALUES ({', '.join(placeholders)})
|
||||
ON CONFLICT ("{primary_key}") DO UPDATE SET
|
||||
{', '.join(f'"{col}" = EXCLUDED."{col}"' for col in columns if col not in [primary_key, 'version', 'server_id'])},
|
||||
version = EXCLUDED.version,
|
||||
server_id = EXCLUDED.server_id
|
||||
WHERE "{table_name}".version < EXCLUDED.version
|
||||
OR ("{table_name}".version = EXCLUDED.version AND "{table_name}".server_id < EXCLUDED.server_id)
|
||||
"""
|
||||
else:
|
||||
# For tables without a primary key, we'll use all columns for conflict resolution
|
||||
insert_query = f"""
|
||||
INSERT INTO "{table_name}" ({', '.join(f'"{col}"' for col in columns)})
|
||||
VALUES ({', '.join(placeholders)})
|
||||
ON CONFLICT DO NOTHING
|
||||
"""
|
||||
|
||||
# debug(f"Generated insert query for {table_name}: {insert_query}")
|
||||
|
||||
affected_rows = 0
|
||||
async for change in tqdm(changes, desc=f"Syncing {table_name}", unit="row"):
|
||||
values = [change[col] for col in columns]
|
||||
# debug(f"Executing query for {table_name} with values: {values}")
|
||||
result = await conn.execute(insert_query, *values)
|
||||
affected_rows += int(result.split()[-1])
|
||||
|
||||
return affected_rows
|
||||
|
||||
except Exception as e:
|
||||
err(f"Error applying batch changes to {table_name}: {str(e)}")
|
||||
err(f"Traceback: {traceback.format_exc()}")
|
||||
return 0
|
||||
|
||||
async def pull_changes(self, source_pool_entry, batch_size=10000):
|
||||
if source_pool_entry['ts_id'] == os.environ.get('TS_ID'):
|
||||
debug("Skipping self-sync")
|
||||
return 0
|
||||
|
||||
total_changes = 0
|
||||
source_id = source_pool_entry['ts_id']
|
||||
source_ip = source_pool_entry['ts_ip']
|
||||
dest_id = os.environ.get('TS_ID')
|
||||
dest_ip = self.local_db['ts_ip']
|
||||
|
||||
info(f"Starting sync from source {source_id} ({source_ip}) to destination {dest_id} ({dest_ip})")
|
||||
|
||||
try:
|
||||
async with self.get_connection(source_pool_entry) as source_conn:
|
||||
async with self.get_connection(self.local_db) as dest_conn:
|
||||
tables = await source_conn.fetch("""
|
||||
SELECT tablename FROM pg_tables
|
||||
WHERE schemaname = 'public'
|
||||
""")
|
||||
|
||||
async for table in tqdm(tables, desc="Syncing tables", unit="table"):
|
||||
table_name = table['tablename']
|
||||
try:
|
||||
if table_name in self.SPECIAL_TABLES:
|
||||
await self.sync_special_table(source_conn, dest_conn, table_name)
|
||||
else:
|
||||
primary_key = await self.ensure_sync_columns(dest_conn, table_name)
|
||||
last_synced_version = await self.get_last_synced_version(dest_conn, table_name, source_id)
|
||||
|
||||
changes = await source_conn.fetch(f"""
|
||||
SELECT * FROM "{table_name}"
|
||||
WHERE version > $1 AND server_id = $2
|
||||
ORDER BY version ASC
|
||||
LIMIT $3
|
||||
""", last_synced_version, source_id, batch_size)
|
||||
|
||||
if changes:
|
||||
changes_count = await self.apply_batch_changes(dest_conn, table_name, changes, primary_key)
|
||||
total_changes += changes_count
|
||||
|
||||
if changes_count > 0:
|
||||
info(f"Synced batch for {table_name}: {changes_count} changes. Total so far: {total_changes}")
|
||||
else:
|
||||
debug(f"No changes to sync for {table_name}")
|
||||
|
||||
except Exception as e:
|
||||
err(f"Error syncing table {table_name}: {str(e)}")
|
||||
err(f"Traceback: {traceback.format_exc()}")
|
||||
|
||||
info(f"Sync complete from {source_id} ({source_ip}) to {dest_id} ({dest_ip}). Total changes: {total_changes}")
|
||||
|
||||
except Exception as e:
|
||||
err(f"Error during sync process: {str(e)}")
|
||||
err(f"Traceback: {traceback.format_exc()}")
|
||||
|
||||
info(f"Sync summary:")
|
||||
info(f" Total changes: {total_changes}")
|
||||
info(f" Tables synced: {len(tables)}")
|
||||
info(f" Source: {source_id} ({source_ip})")
|
||||
info(f" Destination: {dest_id} ({dest_ip})")
|
||||
|
||||
return total_changes
|
||||
|
||||
async def get_online_hosts(self) -> List[Dict[str, Any]]:
|
||||
online_hosts = []
|
||||
for pool_entry in self.POOL:
|
||||
|
|
@ -619,31 +461,6 @@ class APIConfig(BaseModel):
|
|||
err(f"Error checking host {pool_entry['ts_ip']}:{pool_entry['db_port']}: {str(e)}")
|
||||
return online_hosts
|
||||
|
||||
async def push_changes_to_all(self):
|
||||
for pool_entry in self.POOL:
|
||||
if pool_entry['ts_id'] != os.environ.get('TS_ID'):
|
||||
try:
|
||||
await self.push_changes_to_one(pool_entry)
|
||||
except Exception as e:
|
||||
err(f"Error pushing changes to {pool_entry['ts_id']}: {str(e)}")
|
||||
|
||||
|
||||
|
||||
async def get_last_synced_version(self, conn, table_name, server_id):
|
||||
if conn is None:
|
||||
debug(f"Skipping offline server...")
|
||||
return 0
|
||||
|
||||
if table_name in self.SPECIAL_TABLES:
|
||||
debug(f"Skipping get_last_synced_version becaue {table_name} is special.")
|
||||
return 0 # Special handling for tables without version column
|
||||
|
||||
return await conn.fetchval(f"""
|
||||
SELECT COALESCE(MAX(version), 0)
|
||||
FROM "{table_name}"
|
||||
WHERE server_id = $1
|
||||
""", server_id)
|
||||
|
||||
async def check_postgis(self, conn):
|
||||
if conn is None:
|
||||
debug(f"Skipping offline server...")
|
||||
|
|
@ -661,11 +478,134 @@ class APIConfig(BaseModel):
|
|||
err(f"Error checking PostGIS: {str(e)}")
|
||||
return False
|
||||
|
||||
async def sync_special_table(self, source_conn, dest_conn, table_name):
|
||||
async def execute_write_query(self, query: str, *args, table_name: str):
|
||||
if table_name in self.SPECIAL_TABLES:
|
||||
return await self._execute_special_table_write(query, *args, table_name=table_name)
|
||||
|
||||
async with self.get_connection() as conn:
|
||||
if conn is None:
|
||||
raise ConnectionError("Failed to connect to local database")
|
||||
|
||||
# Ensure sync columns exist
|
||||
primary_key = await self.ensure_sync_columns(conn, table_name)
|
||||
|
||||
# Execute the query
|
||||
result = await conn.execute(query, *args)
|
||||
|
||||
# Get the primary key and new version of the affected row
|
||||
if primary_key:
|
||||
affected_row = await conn.fetchrow(f"""
|
||||
SELECT "{primary_key}", version, server_id
|
||||
FROM "{table_name}"
|
||||
WHERE version = (SELECT MAX(version) FROM "{table_name}")
|
||||
""")
|
||||
if affected_row:
|
||||
await self.push_change(table_name, affected_row[primary_key], affected_row['version'], affected_row['server_id'])
|
||||
else:
|
||||
# For tables without a primary key, we'll push all rows
|
||||
await self.push_all_changes(table_name)
|
||||
|
||||
return result
|
||||
|
||||
async def push_change(self, table_name: str, pk_value: Any, version: int, server_id: str):
|
||||
online_hosts = await self.get_online_hosts()
|
||||
for pool_entry in online_hosts:
|
||||
if pool_entry['ts_id'] != os.environ.get('TS_ID'):
|
||||
try:
|
||||
async with self.get_connection(pool_entry) as remote_conn:
|
||||
if remote_conn is None:
|
||||
continue
|
||||
|
||||
# Fetch the updated row from the local database
|
||||
async with self.get_connection() as local_conn:
|
||||
updated_row = await local_conn.fetchrow(f'SELECT * FROM "{table_name}" WHERE "{self.get_primary_key(table_name)}" = $1', pk_value)
|
||||
|
||||
if updated_row:
|
||||
columns = updated_row.keys()
|
||||
placeholders = [f'${i+1}' for i in range(len(columns))]
|
||||
primary_key = self.get_primary_key(table_name)
|
||||
|
||||
insert_query = f"""
|
||||
INSERT INTO "{table_name}" ({', '.join(f'"{col}"' for col in columns)})
|
||||
VALUES ({', '.join(placeholders)})
|
||||
ON CONFLICT ("{primary_key}") DO UPDATE SET
|
||||
{', '.join(f'"{col}" = EXCLUDED."{col}"' for col in columns if col != primary_key)},
|
||||
version = EXCLUDED.version,
|
||||
server_id = EXCLUDED.server_id
|
||||
WHERE "{table_name}".version < EXCLUDED.version
|
||||
OR ("{table_name}".version = EXCLUDED.version AND "{table_name}".server_id < EXCLUDED.server_id)
|
||||
"""
|
||||
await remote_conn.execute(insert_query, *updated_row.values())
|
||||
except Exception as e:
|
||||
err(f"Error pushing change to {pool_entry['ts_id']}: {str(e)}")
|
||||
err(f"Traceback: {traceback.format_exc()}")
|
||||
|
||||
async def push_all_changes(self, table_name: str):
|
||||
online_hosts = await self.get_online_hosts()
|
||||
for pool_entry in online_hosts:
|
||||
if pool_entry['ts_id'] != os.environ.get('TS_ID'):
|
||||
try:
|
||||
async with self.get_connection(pool_entry) as remote_conn:
|
||||
if remote_conn is None:
|
||||
continue
|
||||
|
||||
# Fetch all rows from the local database
|
||||
async with self.get_connection() as local_conn:
|
||||
all_rows = await local_conn.fetch(f'SELECT * FROM "{table_name}"')
|
||||
|
||||
if all_rows:
|
||||
columns = all_rows[0].keys()
|
||||
placeholders = [f'${i+1}' for i in range(len(columns))]
|
||||
|
||||
insert_query = f"""
|
||||
INSERT INTO "{table_name}" ({', '.join(f'"{col}"' for col in columns)})
|
||||
VALUES ({', '.join(placeholders)})
|
||||
ON CONFLICT DO NOTHING
|
||||
"""
|
||||
|
||||
# Use a transaction to insert all rows
|
||||
async with remote_conn.transaction():
|
||||
for row in all_rows:
|
||||
await remote_conn.execute(insert_query, *row.values())
|
||||
except Exception as e:
|
||||
err(f"Error pushing all changes to {pool_entry['ts_id']}: {str(e)}")
|
||||
err(f"Traceback: {traceback.format_exc()}")
|
||||
|
||||
async def _execute_special_table_write(self, query: str, *args, table_name: str):
|
||||
if table_name == 'spatial_ref_sys':
|
||||
return await self.sync_spatial_ref_sys(source_conn, dest_conn)
|
||||
return await self._execute_spatial_ref_sys_write(query, *args)
|
||||
# Add more special cases as needed
|
||||
|
||||
async def _execute_spatial_ref_sys_write(self, query: str, *args):
|
||||
result = None
|
||||
async with self.get_connection() as local_conn:
|
||||
if local_conn is None:
|
||||
raise ConnectionError("Failed to connect to local database")
|
||||
|
||||
# Execute the query locally
|
||||
result = await local_conn.execute(query, *args)
|
||||
|
||||
# Sync the entire spatial_ref_sys table with all online hosts
|
||||
online_hosts = await self.get_online_hosts()
|
||||
for pool_entry in online_hosts:
|
||||
if pool_entry['ts_id'] != os.environ.get('TS_ID'):
|
||||
try:
|
||||
async with self.get_connection(pool_entry) as remote_conn:
|
||||
if remote_conn is None:
|
||||
continue
|
||||
await self.sync_spatial_ref_sys(local_conn, remote_conn)
|
||||
except Exception as e:
|
||||
err(f"Error syncing spatial_ref_sys to {pool_entry['ts_id']}: {str(e)}")
|
||||
err(f"Traceback: {traceback.format_exc()}")
|
||||
|
||||
return result
|
||||
|
||||
def get_primary_key(self, table_name: str) -> str:
|
||||
# This method should return the primary key for the given table
|
||||
# You might want to cache this information for performance
|
||||
# For now, we'll assume it's always 'id', but you should implement proper logic here
|
||||
return 'id'
|
||||
|
||||
async def sync_spatial_ref_sys(self, source_conn, dest_conn):
|
||||
try:
|
||||
# Get all entries from the source
|
||||
|
|
@ -696,7 +636,6 @@ class APIConfig(BaseModel):
|
|||
INSERT INTO spatial_ref_sys ({', '.join(f'"{col}"' for col in columns)})
|
||||
VALUES ({', '.join(placeholders)})
|
||||
"""
|
||||
# debug(f"Inserting new entry for srid {srid}: {insert_query}")
|
||||
await dest_conn.execute(insert_query, *source_entry.values())
|
||||
inserts += 1
|
||||
elif source_entry != dest_dict[srid]:
|
||||
|
|
@ -709,7 +648,6 @@ class APIConfig(BaseModel):
|
|||
proj4text = $4::text
|
||||
WHERE srid = $5::integer
|
||||
"""
|
||||
# debug(f"Updating entry for srid {srid}: {update_query}")
|
||||
await dest_conn.execute(update_query,
|
||||
source_entry['auth_name'],
|
||||
source_entry['auth_srid'],
|
||||
|
|
@ -727,63 +665,6 @@ class APIConfig(BaseModel):
|
|||
err(f"Traceback: {traceback.format_exc()}")
|
||||
return 0
|
||||
|
||||
async def get_most_recent_source(self):
|
||||
most_recent_source = None
|
||||
max_version = -1
|
||||
local_ts_id = os.environ.get('TS_ID')
|
||||
online_hosts = await self.get_online_hosts()
|
||||
num_online_hosts = len(online_hosts)
|
||||
if num_online_hosts > 0:
|
||||
online_ts_ids = [host['ts_id'] for host in online_hosts if host['ts_id'] != local_ts_id]
|
||||
crit(f"Online hosts: {', '.join(online_ts_ids)}")
|
||||
|
||||
for pool_entry in online_hosts:
|
||||
if pool_entry['ts_id'] == local_ts_id:
|
||||
continue # Skip local database
|
||||
|
||||
try:
|
||||
async with self.get_connection(pool_entry) as conn:
|
||||
tables = await conn.fetch("""
|
||||
SELECT tablename FROM pg_tables
|
||||
WHERE schemaname = 'public'
|
||||
""")
|
||||
|
||||
for table in tables:
|
||||
table_name = table['tablename']
|
||||
if table_name in self.SPECIAL_TABLES:
|
||||
continue # Skip special tables for version comparison
|
||||
try:
|
||||
result = await conn.fetchrow(f"""
|
||||
SELECT MAX(version) as max_version, server_id
|
||||
FROM "{table_name}"
|
||||
WHERE version = (SELECT MAX(version) FROM "{table_name}")
|
||||
GROUP BY server_id
|
||||
ORDER BY MAX(version) DESC
|
||||
LIMIT 1
|
||||
""")
|
||||
if result:
|
||||
version, server_id = result['max_version'], result['server_id']
|
||||
info(f"Max version for {pool_entry['ts_id']}, table {table_name}: {version} (from server {server_id})")
|
||||
if version > max_version:
|
||||
max_version = version
|
||||
most_recent_source = pool_entry
|
||||
else:
|
||||
debug(f"No data in table {table_name} for {pool_entry['ts_id']}")
|
||||
except asyncpg.exceptions.UndefinedColumnError:
|
||||
warn(f"Version or server_id column does not exist in table {table_name} for {pool_entry['ts_id']}. Skipping.")
|
||||
except Exception as e:
|
||||
err(f"Error checking version for {pool_entry['ts_id']}, table {table_name}: {str(e)}")
|
||||
|
||||
except asyncpg.exceptions.ConnectionFailureError:
|
||||
warn(f"Failed to connect to database: {pool_entry['ts_ip']}:{pool_entry['db_port']}")
|
||||
except Exception as e:
|
||||
err(f"Unexpected error occurred while checking version for {pool_entry['ts_id']}: {str(e)}")
|
||||
err(f"Traceback: {traceback.format_exc()}")
|
||||
|
||||
return most_recent_source
|
||||
else:
|
||||
warn(f"No other online hosts for sync")
|
||||
return None
|
||||
|
||||
|
||||
class Location(BaseModel):
|
||||
|
|
|
|||
|
|
@ -417,50 +417,49 @@ if API.EXTENSIONS.courtlistener == "on" or API.EXTENSIONS.courtlistener == True:
|
|||
await cl_download_file(download_url, target_path, session)
|
||||
debug(f"Downloaded {file_name} to {target_path}")
|
||||
|
||||
|
||||
@serve.get("/s", response_class=HTMLResponse)
|
||||
async def shortener_form(request: Request):
|
||||
return templates.TemplateResponse("shortener.html", {"request": request})
|
||||
|
||||
@serve.post("/s")
|
||||
async def create_short_url(request: Request, long_url: str = Form(...), custom_code: Optional[str] = Form(None)):
|
||||
async with API.get_connection() as conn:
|
||||
await create_tables(conn)
|
||||
await create_tables()
|
||||
|
||||
if custom_code:
|
||||
if len(custom_code) != 3 or not custom_code.isalnum():
|
||||
return templates.TemplateResponse("shortener.html", {"request": request, "error": "Custom code must be 3 alphanumeric characters"})
|
||||
if custom_code:
|
||||
if len(custom_code) != 3 or not custom_code.isalnum():
|
||||
return templates.TemplateResponse("shortener.html", {"request": request, "error": "Custom code must be 3 alphanumeric characters"})
|
||||
|
||||
existing = await conn.fetchval('SELECT 1 FROM short_urls WHERE short_code = $1', custom_code)
|
||||
if existing:
|
||||
return templates.TemplateResponse("shortener.html", {"request": request, "error": "Custom code already in use"})
|
||||
existing = await API.execute_write_query('SELECT 1 FROM short_urls WHERE short_code = $1', custom_code, table_name="short_urls")
|
||||
if existing:
|
||||
return templates.TemplateResponse("shortener.html", {"request": request, "error": "Custom code already in use"})
|
||||
|
||||
short_code = custom_code
|
||||
else:
|
||||
chars = string.ascii_letters + string.digits
|
||||
while True:
|
||||
short_code = ''.join(random.choice(chars) for _ in range(3))
|
||||
existing = await conn.fetchval('SELECT 1 FROM short_urls WHERE short_code = $1', short_code)
|
||||
if not existing:
|
||||
break
|
||||
short_code = custom_code
|
||||
else:
|
||||
chars = string.ascii_letters + string.digits
|
||||
while True:
|
||||
short_code = ''.join(random.choice(chars) for _ in range(3))
|
||||
existing = await API.execute_write_query('SELECT 1 FROM short_urls WHERE short_code = $1', short_code, table_name="short_urls")
|
||||
if not existing:
|
||||
break
|
||||
|
||||
await conn.execute(
|
||||
'INSERT INTO short_urls (short_code, long_url) VALUES ($1, $2)',
|
||||
short_code, long_url
|
||||
)
|
||||
await API.execute_write_query(
|
||||
'INSERT INTO short_urls (short_code, long_url) VALUES ($1, $2)',
|
||||
short_code, long_url,
|
||||
table_name="short_urls"
|
||||
)
|
||||
|
||||
short_url = f"https://sij.ai/{short_code}"
|
||||
return templates.TemplateResponse("shortener.html", {"request": request, "short_url": short_url})
|
||||
|
||||
async def create_tables(conn):
|
||||
await conn.execute('''
|
||||
async def create_tables():
|
||||
await API.execute_write_query('''
|
||||
CREATE TABLE IF NOT EXISTS short_urls (
|
||||
short_code VARCHAR(3) PRIMARY KEY,
|
||||
long_url TEXT NOT NULL,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
''')
|
||||
await conn.execute('''
|
||||
''', table_name="short_urls")
|
||||
await API.execute_write_query('''
|
||||
CREATE TABLE IF NOT EXISTS click_logs (
|
||||
id SERIAL PRIMARY KEY,
|
||||
short_code VARCHAR(3) REFERENCES short_urls(short_code),
|
||||
|
|
@ -468,47 +467,50 @@ async def create_tables(conn):
|
|||
ip_address TEXT,
|
||||
user_agent TEXT
|
||||
)
|
||||
''')
|
||||
''', table_name="click_logs")
|
||||
|
||||
@serve.get("/{short_code}", response_class=RedirectResponse, status_code=301)
|
||||
async def redirect_short_url(request: Request, short_code: str = PathParam(..., min_length=3, max_length=3)):
|
||||
if request.headers.get('host') != 'sij.ai':
|
||||
raise HTTPException(status_code=404, detail="Not Found")
|
||||
|
||||
async with API.get_connection() as conn:
|
||||
result = await conn.fetchrow(
|
||||
'SELECT long_url FROM short_urls WHERE short_code = $1',
|
||||
short_code
|
||||
)
|
||||
result = await API.execute_write_query(
|
||||
'SELECT long_url FROM short_urls WHERE short_code = $1',
|
||||
short_code,
|
||||
table_name="short_urls"
|
||||
)
|
||||
|
||||
if result:
|
||||
await conn.execute(
|
||||
'INSERT INTO click_logs (short_code, ip_address, user_agent) VALUES ($1, $2, $3)',
|
||||
short_code, request.client.host, request.headers.get("user-agent")
|
||||
)
|
||||
return result['long_url']
|
||||
else:
|
||||
raise HTTPException(status_code=404, detail="Short URL not found")
|
||||
if result:
|
||||
await API.execute_write_query(
|
||||
'INSERT INTO click_logs (short_code, ip_address, user_agent) VALUES ($1, $2, $3)',
|
||||
short_code, request.client.host, request.headers.get("user-agent"),
|
||||
table_name="click_logs"
|
||||
)
|
||||
return result['long_url']
|
||||
else:
|
||||
raise HTTPException(status_code=404, detail="Short URL not found")
|
||||
|
||||
@serve.get("/analytics/{short_code}")
|
||||
async def get_analytics(short_code: str):
|
||||
async with API.get_connection() as conn:
|
||||
url_info = await conn.fetchrow(
|
||||
'SELECT long_url, created_at FROM short_urls WHERE short_code = $1',
|
||||
short_code
|
||||
)
|
||||
if not url_info:
|
||||
raise HTTPException(status_code=404, detail="Short URL not found")
|
||||
url_info = await API.execute_write_query(
|
||||
'SELECT long_url, created_at FROM short_urls WHERE short_code = $1',
|
||||
short_code,
|
||||
table_name="short_urls"
|
||||
)
|
||||
if not url_info:
|
||||
raise HTTPException(status_code=404, detail="Short URL not found")
|
||||
|
||||
click_count = await conn.fetchval(
|
||||
'SELECT COUNT(*) FROM click_logs WHERE short_code = $1',
|
||||
short_code
|
||||
)
|
||||
click_count = await API.execute_write_query(
|
||||
'SELECT COUNT(*) FROM click_logs WHERE short_code = $1',
|
||||
short_code,
|
||||
table_name="click_logs"
|
||||
)
|
||||
|
||||
clicks = await conn.fetch(
|
||||
'SELECT clicked_at, ip_address, user_agent FROM click_logs WHERE short_code = $1 ORDER BY clicked_at DESC LIMIT 100',
|
||||
short_code
|
||||
)
|
||||
clicks = await API.execute_write_query(
|
||||
'SELECT clicked_at, ip_address, user_agent FROM click_logs WHERE short_code = $1 ORDER BY clicked_at DESC LIMIT 100',
|
||||
short_code,
|
||||
table_name="click_logs"
|
||||
)
|
||||
|
||||
return {
|
||||
"short_code": short_code,
|
||||
|
|
@ -520,6 +522,7 @@ async def get_analytics(short_code: str):
|
|||
|
||||
|
||||
|
||||
|
||||
async def forward_traffic(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, destination: str):
|
||||
dest_host, dest_port = destination.split(':')
|
||||
dest_port = int(dest_port)
|
||||
|
|
|
|||
Loading…
Reference in a new issue