yearning-cli 0.2.0__tar.gz → 0.2.1__tar.gz

{yearning_cli-0.2.0 → yearning_cli-0.2.1}/.github/workflows/docker-image.yml

@@ -17,7 +17,7 @@ jobs:
 
     steps:
       - name: Checkout
-        uses: actions/checkout@v6
+        uses: actions/checkout@v7
 
       - name: Set up QEMU
         uses: docker/setup-qemu-action@v4

{yearning_cli-0.2.0 → yearning_cli-0.2.1}/.github/workflows/publish-pypi.yml

@@ -15,7 +15,7 @@ jobs:
 
     steps:
       - name: Checkout
-        uses: actions/checkout@v6
+        uses: actions/checkout@v7
 
       - name: Set up Python
         uses: actions/setup-python@v6

{yearning_cli-0.2.0 → yearning_cli-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: yearning-cli
-Version: 0.2.0
+Version: 0.2.1
 Summary: Yearning MySQL Audit Platform CLI Tool
 Project-URL: Homepage, https://github.com/la3rence/yearning-cli
 Project-URL: Repository, https://github.com/la3rence/yearning-cli
@@ -148,6 +148,12 @@ mysql -h 127.0.0.1 -P 3307 -u 数据源名 数据库名
 # Host: 127.0.0.1 / 机器局域网 IP / 反代入口  Port: 3307  Username: 数据源名称  Password: 留空
 ```
 
+代理会根据数据库名自动切换 Yearning 数据源：
+
+- 执行 `USE app_billing` 或客户端发送 COM_INIT_DB 时，代理会在可用数据源中查找包含该库的数据源并切换
+- 查询中出现 `db.table` 形式（例如 `SELECT count(1) FROM app_billing.invoice`）时，也会先按 `db` 解析数据源
+- 如果同名数据库存在于多个数据源，代理会优先选择数据源名称与库名更匹配的源（例如 `app_billing` 优先匹配名称包含 `billing` 的数据源）；无法判断时保留当前数据源
+
 代理参数：
 
 | 参数 | 说明 | 默认值 |

{yearning_cli-0.2.0 → yearning_cli-0.2.1}/README.md

@@ -133,6 +133,12 @@ mysql -h 127.0.0.1 -P 3307 -u 数据源名 数据库名
 # Host: 127.0.0.1 / 机器局域网 IP / 反代入口  Port: 3307  Username: 数据源名称  Password: 留空
 ```
 
+代理会根据数据库名自动切换 Yearning 数据源：
+
+- 执行 `USE app_billing` 或客户端发送 COM_INIT_DB 时，代理会在可用数据源中查找包含该库的数据源并切换
+- 查询中出现 `db.table` 形式（例如 `SELECT count(1) FROM app_billing.invoice`）时，也会先按 `db` 解析数据源
+- 如果同名数据库存在于多个数据源，代理会优先选择数据源名称与库名更匹配的源（例如 `app_billing` 优先匹配名称包含 `billing` 的数据源）；无法判断时保留当前数据源
+
 代理参数：
 
 | 参数 | 说明 | 默认值 |

{yearning_cli-0.2.0 → yearning_cli-0.2.1}/mysql_proxy/README.md

@@ -26,6 +26,12 @@ MySQL 协议代理模块，让 DBeaver、Navicat、MySQL CLI 等客户端通过
 
 密码字段无需填写，任意值均可。
 
+连接后也支持按数据库名自动切换数据源：
+
+- `USE dbname` 和客户端的切换数据库命令会查找包含该库的数据源
+- 查询中的 `dbname.table_name` 会按 `dbname` 自动路由到对应数据源
+- 同名数据库出现在多个数据源时，会优先选择数据源名称与库名更匹配的源；无法判断时继续使用当前数据源
+
 ## 前提条件
 
 1. 已运行 `sql login` 登录 Yearning，或已配置 `YEARNING_USER`/`YEARNING_PASS`、`~/.sqlrc` 的 `yearning_user`/`yearning_pass` 供 token 失效时自动登录

{yearning_cli-0.2.0 → yearning_cli-0.2.1}/mysql_proxy/server.py

@@ -4,6 +4,7 @@ import hashlib
 import logging
 import os
 import random
+import re
 import socket
 import sys
 import threading
@@ -42,9 +43,22 @@ class ConnectionHandler:
     _lock = threading.Lock()
     _cache_lock = threading.Lock()
     _source_cache = {}
+    _database_source_cache = {}
     _query_order_cache = {}
     SOURCE_CACHE_TTL = 60
+    DATABASE_SOURCE_CACHE_TTL = 300
     QUERY_ORDER_TTL = 600
+    _QUALIFIED_TABLE_RE = re.compile(
+        r"\b(?:FROM|JOIN|UPDATE|INTO|TABLE|DESC|DESCRIBE)\s+"
+        r"`?([A-Za-z_][\w$]*)`?\s*\.\s*`?([A-Za-z_][\w$]*)`?",
+        re.IGNORECASE,
+    )
+    _SYSTEM_SCHEMAS = {
+        "information_schema",
+        "mysql",
+        "performance_schema",
+        "sys",
+    }
 
     def __init__(self, conn: socket.socket, addr, config: dict):
         self.conn = conn
@@ -202,7 +216,9 @@ class ConnectionHandler:
             self._init_yearning()
         return self.yearning and self.yearning.token
 
-    def _ensure_query_order(self):
+    def _ensure_query_order(self, source_id=None):
+        if source_id:
+            self.source_id = source_id
         if not self._ensure_yearning_initialized() or not self.source_id:
             return False
 
@@ -316,8 +332,7 @@ class ConnectionHandler:
                 self._handle_query(sql)
             elif cmd == 0x02:  # COM_INIT_DB (USE database)
                 db = payload[1:].decode("utf-8", errors="replace").strip()
-                self.current_db = db
-                self._send(make_ok(message=f"Changed to database '{db}'"))
+                self._handle_use_database(db)
             elif cmd == 0x0E:  # COM_PING
                 self._send(make_ok())
             elif cmd == 0x01:  # COM_QUIT
@@ -358,9 +373,11 @@ class ConnectionHandler:
             self._send(make_ok(message="OK"))
             return
         if upper.startswith("USE "):
-            db = sql[3:].strip().strip("`\"'")
-            self.current_db = db
-            self._send(make_ok(message=f"Changed to database '{db}'"))
+            db = self._parse_use_database(sql)
+            if not db:
+                self._send(make_err(1064, "Syntax error"))
+                return
+            self._handle_use_database(db)
             return
         # SELECT @@version, SELECT @@version_comment, etc.
         if upper.startswith("SELECT @@"):
@@ -404,13 +421,14 @@ class ConnectionHandler:
             self._send(make_err(1046, "No data source configured. Set source_id in config."))
             return
 
-        if not self.current_db:
+        source_id, database = self._resolve_query_route(sql)
+        if not database:
             self._send(make_err(1046, "No database selected"))
             return
 
-        self._ensure_query_order()
+        self._ensure_query_order(source_id)
         try:
-            result = self.yearning.execute_query(self.source_id, self.current_db, sql, timeout=120)
+            result = self.yearning.execute_query(source_id, database, sql, timeout=120)
         except Exception as e:
             logger.error("Query error: %s", e)
             self._send(make_err(1105, str(e), "HY000"))
@@ -422,9 +440,9 @@ class ConnectionHandler:
             if "工单" in error_msg or "expired" in error_msg.lower():
                 logger.info("Query order expired, recreating...")
                 try:
-                    if self.yearning.create_query_order(self.source_id, verify_wait=0.1):
+                    if self.yearning.create_query_order(source_id, verify_wait=0.1):
                         self._mark_query_order_ready()
-                    result = self.yearning.execute_query(self.source_id, self.current_db, sql, timeout=120)
+                    result = self.yearning.execute_query(source_id, database, sql, timeout=120)
                     if result["success"]:
                         self._mark_query_order_ready()
                         self._send_result(result)
@@ -437,6 +455,169 @@ class ConnectionHandler:
         self._mark_query_order_ready()
         self._send_result(result)
 
+    @staticmethod
+    def _parse_use_database(sql: str) -> str:
+        """Return database from a USE statement."""
+        value = sql.strip().rstrip(";").strip()
+        if not value[:3].upper() == "USE":
+            return ""
+        db = value[3:].strip()
+        if not db:
+            return ""
+        return db.strip("`\"'")
+
+    def _handle_use_database(self, db: str):
+        """Switch the active database and source when the database is known elsewhere."""
+        db = db.strip().rstrip(";").strip("`\"'")
+        if not db:
+            self._send(make_err(1064, "Syntax error"))
+            return
+        self._ensure_yearning_initialized()
+        if self.yearning and self.yearning.token:
+            source_id = self._resolve_source_for_database(db)
+            if source_id:
+                self._set_route(source_id, db)
+        self.current_db = db
+        self._send(make_ok(message=f"Changed to database '{db}'"))
+
+    def _set_route(self, source_id: str, database: str):
+        """Update the connection route when it changes."""
+        if source_id and source_id != self.source_id:
+            logger.info(
+                "Switching source for database '%s': %s -> %s",
+                database,
+                self.source_id[:12] if self.source_id else "NONE",
+                source_id[:12],
+            )
+            self.source_id = source_id
+        self.current_db = database
+
+    def _resolve_query_route(self, sql: str):
+        """Resolve the source/database for a query before sending it to Yearning."""
+        explicit_db = self._first_explicit_database(sql)
+        database = explicit_db or self.current_db
+        if database:
+            source_id = self._resolve_source_for_database(database)
+            if source_id:
+                self._set_route(source_id, database)
+        return self.source_id, database
+
+    def _first_explicit_database(self, sql: str) -> str:
+        """Find the first non-system schema in a qualified table reference."""
+        for match in self._QUALIFIED_TABLE_RE.finditer(sql):
+            schema = match.group(1)
+            if schema.lower() not in self._SYSTEM_SCHEMAS:
+                return schema
+        return ""
+
+    def _resolve_source_for_database(self, database: str):
+        """Find the Yearning source that contains a database."""
+        if not database or not self._ensure_yearning_initialized():
+            return ""
+
+        key = self._cache_key("database_source", database)
+        now = time.monotonic()
+        with self._cache_lock:
+            entry = self._database_source_cache.get(key)
+            if entry and entry["expires_at"] > now:
+                return entry["source_id"]
+
+        matches = []
+        for source in self._list_sources_cached():
+            source_id = source.get("source_id", "")
+            if not source_id:
+                continue
+            if self._source_has_database(source_id, database):
+                matches.append(source)
+
+        source_id = ""
+        if len(matches) == 1:
+            source_id = matches[0].get("source_id", "")
+        elif len(matches) > 1:
+            source_id = self._choose_source_for_database(database, matches)
+        if len(matches) > 1 and not source_id:
+            logger.warning(
+                "Database '%s' exists in multiple sources; keeping current source",
+                database,
+            )
+        self._cache_database_source(database, source_id)
+        return source_id
+
+    def _choose_source_for_database(self, database: str, sources: list):
+        """Choose the most specific source when a database exists in multiple sources."""
+        ranked = [
+            (self._database_source_score(database, source), source)
+            for source in sources
+        ]
+        ranked = [item for item in ranked if item[0] > 0]
+        if not ranked:
+            return ""
+        ranked.sort(key=lambda item: item[0], reverse=True)
+        best_score, best_source = ranked[0]
+        if len(ranked) > 1 and ranked[1][0] == best_score:
+            logger.warning(
+                "Database '%s' has equally specific source matches: %s",
+                database,
+                ", ".join(
+                    source.get("source", "")
+                    for score, source in ranked
+                    if score == best_score
+                ),
+            )
+            return ""
+        source_id = best_source.get("source_id", "")
+        logger.info(
+            "Resolved database '%s' to source '%s' by source-name match",
+            database,
+            best_source.get("source", source_id),
+        )
+        return source_id
+
+    @classmethod
+    def _database_source_score(cls, database: str, source: dict) -> int:
+        source_name = source.get("source", "")
+        db_norm = cls._normalize_identifier(database)
+        source_norm = cls._normalize_identifier(source_name)
+        if not db_norm or not source_norm:
+            return 0
+        if source_norm == db_norm:
+            return 100
+        if db_norm in source_norm:
+            return 90
+
+        db_tokens = [token for token in db_norm.split("_") if token]
+        source_tokens = set(token for token in source_norm.split("_") if token)
+        score = 0
+        if db_tokens and db_tokens[-1] in source_tokens:
+            score += 70
+        score += 10 * len(set(db_tokens) & source_tokens)
+        return score
+
+    @staticmethod
+    def _normalize_identifier(value: str) -> str:
+        return re.sub(r"_+", "_", re.sub(r"[^a-z0-9]+", "_", value.lower())).strip("_")
+
+    def _source_has_database(self, source_id: str, database: str) -> bool:
+        try:
+            return database in self.yearning.list_databases(source_id)
+        except Exception as e:
+            logger.debug(
+                "Failed to list databases for source %s: %s",
+                source_id[:12] if source_id else "NONE",
+                e,
+            )
+            return False
+
+    def _cache_database_source(self, database: str, source_id: str):
+        if not source_id:
+            return
+        key = self._cache_key("database_source", database)
+        with self._cache_lock:
+            self._database_source_cache[key] = {
+                "source_id": source_id,
+                "expires_at": time.monotonic() + self.DATABASE_SOURCE_CACHE_TTL,
+            }
+
     def _fake_sysvar(self, var: str) -> str:
         """Return fake system variable values for client compatibility."""
         var = var.upper().lstrip("@")
@@ -530,12 +711,13 @@ class ConnectionHandler:
         if not self.source_id:
             self._send(make_err(1046, "No data source configured"))
             return
-        if not self.current_db:
+        source_id, database = self._resolve_query_route(sql)
+        if not database:
             self._send(make_err(1046, "No database selected"))
             return
-        self._ensure_query_order()
+        self._ensure_query_order(source_id)
         try:
-            result = self.yearning.execute_query(self.source_id, self.current_db, sql, timeout=120)
+            result = self.yearning.execute_query(source_id, database, sql, timeout=120)
         except Exception as e:
             self._send(make_err(1105, str(e)))
             return

{yearning_cli-0.2.0 → yearning_cli-0.2.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "yearning-cli"
-version = "0.2.0"
+version = "0.2.1"
 description = "Yearning MySQL Audit Platform CLI Tool"
 readme = "README.md"
 authors = [

yearning_cli-0.2.1/tests/test_proxy_routing.py

@@ -0,0 +1,98 @@
+import unittest
+
+from mysql_proxy.server import ConnectionHandler
+
+
+class FakeYearning:
+    token = "token"
+
+    def __init__(self):
+        self.databases = {
+            "source-shared": ["app_billing", "app_orders"],
+            "source-billing": ["app_billing"],
+            "source-orders": ["app_orders"],
+            "source-a": ["shared"],
+            "source-b": ["shared"],
+        }
+
+    def list_sources(self):
+        return [
+            {"source": "shared-apps", "source_id": "source-shared"},
+            {"source": "prod-billing", "source_id": "source-billing"},
+            {"source": "readonly-orders", "source_id": "source-orders"},
+            {"source": "a", "source_id": "source-a"},
+            {"source": "b", "source_id": "source-b"},
+        ]
+
+    def list_databases(self, source_id):
+        return self.databases[source_id]
+
+
+def make_handler():
+    handler = ConnectionHandler.__new__(ConnectionHandler)
+    handler.config = {"base_url": "http://yearning.local"}
+    handler.yearning = FakeYearning()
+    handler.source_id = "source-orders"
+    handler.current_db = "app_orders"
+    return handler
+
+
+class ProxyRoutingTest(unittest.TestCase):
+    def test_explicit_database_routes_to_matching_source(self):
+        handler = make_handler()
+
+        source_id, database = handler._resolve_query_route(
+            "select count(1) from app_orders.order_record"
+        )
+
+        self.assertEqual(source_id, "source-orders")
+        self.assertEqual(database, "app_orders")
+        self.assertEqual(handler.source_id, "source-orders")
+        self.assertEqual(handler.current_db, "app_orders")
+
+    def test_specific_source_beats_broad_source_when_database_is_shared(self):
+        handler = make_handler()
+
+        source_id, database = handler._resolve_query_route(
+            "select * from app_billing.invoice_config"
+        )
+
+        self.assertEqual(source_id, "source-billing")
+        self.assertEqual(database, "app_billing")
+        self.assertEqual(handler.source_id, "source-billing")
+        self.assertEqual(handler.current_db, "app_billing")
+
+    def test_current_database_routes_before_unqualified_query(self):
+        handler = make_handler()
+        handler.current_db = "app_orders"
+
+        source_id, database = handler._resolve_query_route(
+            "select count(1) from order_record"
+        )
+
+        self.assertEqual(source_id, "source-orders")
+        self.assertEqual(database, "app_orders")
+
+    def test_column_reference_is_not_treated_as_database(self):
+        handler = make_handler()
+
+        source_id, database = handler._resolve_query_route(
+            "select t.id from order_record t"
+        )
+
+        self.assertEqual(source_id, "source-orders")
+        self.assertEqual(database, "app_orders")
+
+    def test_ambiguous_database_keeps_current_source(self):
+        handler = make_handler()
+
+        source_id, database = handler._resolve_query_route(
+            "select count(1) from shared.some_table"
+        )
+
+        self.assertEqual(source_id, "source-orders")
+        self.assertEqual(database, "shared")
+
+
+if __name__ == "__main__":
+    unittest.main()

{yearning_cli-0.2.0 → yearning_cli-0.2.1}/uv.lock

@@ -218,7 +218,7 @@ wheels = [
 
 [[package]]
 name = "yearning-cli"
-version = "0.2.0"
+version = "0.2.1"
 source = { editable = "." }
 dependencies = [
     { name = "msgpack" },