code-node-graph 0.1.0__py3-none-any.whl

code_node/libraries/mongo.py

@@ -0,0 +1,169 @@
+"""
+MongoDB through PyMongo (and Motor, its async twin): a client, its databases and their collections, by attribute
+or by name, and what each collection method does.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from . import Call, Handle, Library, Use
+
+_CLIENTS = ("MongoClient", "AsyncMongoClient", "AsyncIOMotorClient", "MotorClient")
+
+_CLIENT_METHODS = frozenset(
+    {
+        "get_database",
+        "get_default_database",
+        "list_database_names",
+        "list_databases",
+        "drop_database",
+        "close",
+        "start_session",
+        "server_info",
+        "watch",
+        "address",
+        "options",
+        "admin",
+        "bulk_write",
+        "with_options",
+    }
+)
+
+_DATABASE_METHODS = frozenset(
+    {
+        "get_collection",
+        "list_collection_names",
+        "list_collections",
+        "create_collection",
+        "drop_collection",
+        "command",
+        "aggregate",
+        "watch",
+        "client",
+        "name",
+        "with_options",
+        "validate_collection",
+        "dereference",
+        "cursor_command",
+    }
+)
+
+_READS = frozenset(
+    {
+        "find",
+        "find_one",
+        "count_documents",
+        "estimated_document_count",
+        "distinct",
+        "aggregate",
+        "watch",
+        "find_raw_batches",
+        "aggregate_raw_batches",
+        "list_indexes",
+        "index_information",
+        "options",
+    }
+)
+
+_WRITES = frozenset(
+    {
+        "insert_one",
+        "insert_many",
+        "update_one",
+        "update_many",
+        "replace_one",
+        "delete_one",
+        "delete_many",
+        "find_one_and_update",
+        "find_one_and_replace",
+        "find_one_and_delete",
+        "bulk_write",
+        "create_index",
+        "create_indexes",
+        "drop_index",
+        "drop_indexes",
+        "drop",
+        "rename",
+    }
+)
+
+_COLLECTION_METHODS = _READS | _WRITES | {"with_options", "name", "full_name", "database"}
+
+_MANY = {
+    "insert_one": "insert_many()",
+    "update_one": "bulk_write() with UpdateOne",
+    "delete_one": "delete_many()",
+    "replace_one": "bulk_write() with ReplaceOne",
+    "find_one": "one find() with $in",
+}
+
+
+class MongoLibrary(Library):
+    """
+    PyMongo and Motor: MongoClient(...), client.shop.orders (or client['shop']['orders']), and its methods.
+    """
+
+    name = "mongodb"
+    title = "MongoDB"
+    modules = ("pymongo", "motor")
+    icon = "database"
+
+    def make(self, name: str, call: Call) -> Handle | Use | None:
+        if name.rsplit(".", 1)[-1] in _CLIENTS:
+            return Handle("client", "MongoDB", opens=True)
+        return None
+
+    def attribute(self, handle: Handle, name: str) -> Handle | None:
+        if name.startswith("_"):
+            return None
+        if handle.kind == "client" and name not in _CLIENT_METHODS:
+            return handle.but("database", database=name)
+        if handle.kind == "database" and name not in _DATABASE_METHODS:
+            return handle.but("collection", collection=name)
+        if handle.kind == "collection" and name not in _COLLECTION_METHODS:  # db.orders.archive: a sub-collection
+            return handle.but("collection", collection=f"{handle.fact('collection')}.{name}")
+        return None
+
+    def item(self, handle: Handle, key: Any) -> Handle | None:
+        name = key if isinstance(key, str) else "?"
+        if handle.kind == "client":
+            return handle.but("database", database=name)
+        if handle.kind == "database":
+            return handle.but("collection", collection=name)
+        return None
+
+    def call(self, handle: Handle, method: str, call: Call) -> Handle | Use | None:
+        if handle.kind == "client" and method in ("get_database", "get_default_database"):
+            return handle.but("database", database=str(call.value(0, "name") or "?"))
+        if handle.kind == "database" and method in ("get_collection", "create_collection"):
+            return handle.but("collection", collection=str(call.value(0, "name") or "?"))
+        if handle.kind != "collection":
+            return None
+        collection = handle.fact("collection")
+        flag = why = ""
+        where = call.inputs(0, "filter")
+        if (
+            method in ("find", "find_one", "delete_many", "update_many", "count_documents")
+            and where
+            and "$where" in call.code(0, "filter", most=400)
+        ):
+            flag, why = (
+                "$where from the input",
+                (
+                    f"runs JavaScript built from {where[0]} on the server ($where): whoever sends it can run their own. "
+                    "Match the fields with a plain filter instead."
+                ),
+            )
+        if method in _READS:
+            return Use(method, "reads", collection, flag=flag, why=why, batch=_MANY.get(method, ""))
+        if method in _WRITES:
+            return Use(method, "writes", collection, flag=flag, why=why, batch=_MANY.get(method, ""))
+        return None
+
+    def describe(self, handle: Handle) -> str:
+        if handle.kind == "database":
+            return f"the MongoDB database {handle.fact('database')}"
+        if handle.kind == "collection":
+            return f"the MongoDB collection {handle.fact('collection')}"
+        return "a MongoDB client"

code_node/libraries/redis.py

@@ -0,0 +1,253 @@
+"""
+Redis through redis-py: the client (or a cluster's), its pipelines and pub/sub, and what each command does.
+"""
+
+from __future__ import annotations
+
+from . import Call, Handle, Library, Use
+
+_CLIENTS = ("Redis", "StrictRedis", "RedisCluster", "from_url")
+
+_READS = frozenset(
+    {
+        "get",
+        "mget",
+        "getrange",
+        "strlen",
+        "exists",
+        "type",
+        "ttl",
+        "pttl",
+        "hget",
+        "hgetall",
+        "hmget",
+        "hkeys",
+        "hvals",
+        "hlen",
+        "hexists",
+        "hscan",
+        "hscan_iter",
+        "lrange",
+        "llen",
+        "lindex",
+        "lpos",
+        "smembers",
+        "sismember",
+        "smismember",
+        "scard",
+        "srandmember",
+        "sscan",
+        "sscan_iter",
+        "sinter",
+        "sunion",
+        "sdiff",
+        "zrange",
+        "zrevrange",
+        "zrangebyscore",
+        "zrevrangebyscore",
+        "zscore",
+        "zmscore",
+        "zcard",
+        "zcount",
+        "zrank",
+        "zrevrank",
+        "zscan",
+        "zscan_iter",
+        "scan",
+        "scan_iter",
+        "keys",
+        "xread",
+        "xreadgroup",
+        "xrange",
+        "xrevrange",
+        "xlen",
+        "blpop",
+        "brpop",
+        "brpoplpush",
+        "blmove",
+        "bzpopmin",
+        "bzpopmax",
+        "pfcount",
+        "getbit",
+        "bitcount",
+        "geopos",
+        "geodist",
+        "geosearch",
+        "dbsize",
+        "json",
+        "ping",
+        "info",
+    }
+)
+
+_WRITES = frozenset(
+    {
+        "set",
+        "setex",
+        "psetex",
+        "setnx",
+        "mset",
+        "msetnx",
+        "getset",
+        "getdel",
+        "getex",
+        "delete",
+        "unlink",
+        "expire",
+        "pexpire",
+        "expireat",
+        "pexpireat",
+        "persist",
+        "rename",
+        "renamenx",
+        "incr",
+        "incrby",
+        "incrbyfloat",
+        "decr",
+        "decrby",
+        "append",
+        "setrange",
+        "hset",
+        "hmset",
+        "hsetnx",
+        "hdel",
+        "hincrby",
+        "hincrbyfloat",
+        "lpush",
+        "rpush",
+        "lpushx",
+        "rpushx",
+        "lpop",
+        "rpop",
+        "lrem",
+        "lset",
+        "ltrim",
+        "linsert",
+        "lmove",
+        "rpoplpush",
+        "sadd",
+        "srem",
+        "spop",
+        "smove",
+        "sinterstore",
+        "sunionstore",
+        "sdiffstore",
+        "zadd",
+        "zrem",
+        "zincrby",
+        "zpopmin",
+        "zpopmax",
+        "zremrangebyscore",
+        "zremrangebyrank",
+        "xack",
+        "xdel",
+        "xtrim",
+        "xgroup_create",
+        "pfadd",
+        "setbit",
+        "geoadd",
+        "copy",
+        "flushdb",
+        "flushall",
+    }
+)
+
+_SENDS = frozenset({"publish", "spublish", "xadd"})
+
+_RUNS = frozenset({"eval", "evalsha", "fcall", "script_load", "register_script"})
+
+_MANY = {"get": "mget()", "set": "mset()", "hget": "hmget()", "delete": "one delete() with every key"}
+
+
+class RedisLibrary(Library):
+    """
+    redis-py: redis.Redis(...) and redis.from_url(...), its pipelines and pub/sub.
+    """
+
+    name = "redis"
+    title = "Redis"
+    modules = ("redis",)
+    icon = "database"
+
+    def make(self, name: str, call: Call) -> Handle | Use | None:
+        last = name.rsplit(".", 1)[-1]
+        if last in _CLIENTS:
+            given = call.given("socket_timeout") or "socket_timeout=" in call.code(0, most=200)
+            return Handle(
+                "client",
+                "Redis",
+                opens=True,
+                waits="" if given else "as long as Redis takes",
+                fix=f"{name.rsplit('.', 1)[-1]}(..., socket_timeout=5)",
+            )
+        if last in ("ConnectionPool", "BlockingConnectionPool"):
+            given = call.given("socket_timeout")
+            return Handle(
+                "pool", "Redis", waits="" if given else "as long as Redis takes", fix=f"{last}(..., socket_timeout=5)"
+            )
+        return None
+
+    def call(self, handle: Handle, method: str, call: Call) -> Handle | Use | None:
+        if method == "pipeline":
+            return handle.but("pipeline")
+        if method == "pubsub":
+            return handle.but("pubsub")
+        if method == "lock":
+            return Use("LOCK", "writes", call.shown(0, "name"), gives=handle.but("lock"))
+        if handle.kind == "pubsub":
+            if method in ("subscribe", "psubscribe"):
+                return Use(method.upper(), "reads", call.shown(0))
+            if method in ("listen", "get_message"):
+                return None
+        if handle.kind == "lock" or method in ("execute", "close", "reset", "watch", "multi"):
+            return None
+        target = call.shown(0, "name") if call.node.args or call.given("name", "keys", "channel") else ""
+        operation = {
+            "scan_iter": "SCAN",
+            "hscan_iter": "HSCAN",
+            "sscan_iter": "SSCAN",
+            "zscan_iter": "ZSCAN",
+            "register_script": "SCRIPT LOAD",
+            "script_load": "SCRIPT LOAD",
+        }.get(method, method.upper())
+        batch = "" if handle.kind == "pipeline" else (_MANY.get(method, "") or "a pipeline()")
+        if method in _READS:
+            if method == "keys":
+                return Use(
+                    operation,
+                    "reads",
+                    target,
+                    flag="KEYS blocks Redis",
+                    batch=batch,
+                    why=(
+                        "walks every key in the database, and Redis answers nothing else until it's done: on a big "
+                        "database that stalls every client. Use scan_iter(match=...) instead."
+                    ),
+                )
+            return Use(operation, "reads", target, batch=batch)
+        if method in _WRITES:
+            if method in ("flushdb", "flushall"):
+                return Use(
+                    operation,
+                    "writes",
+                    flag="empties the database",
+                    why=(
+                        f"deletes every key{' in every database' if method == 'flushall' else ''}: make sure this can only "
+                        "run where that's meant."
+                    ),
+                )
+            return Use(operation, "writes", target, batch=batch)
+        if method in _SENDS:
+            return Use(operation, "sends", target, icon="queue", batch=batch if method == "xadd" else "")
+        if method in _RUNS:
+            return Use("EVAL" if method == "eval" else operation, "runs", "a Lua script")
+        return None
+
+    def describe(self, handle: Handle) -> str:
+        return {
+            "client": "a Redis client",
+            "pool": "a Redis connection pool",
+            "pipeline": "a Redis pipeline",
+            "pubsub": "a Redis pub/sub",
+            "lock": "a Redis lock",
+        }.get(handle.kind, "a Redis " + handle.kind)

code_node/libraries/sql.py

@@ -0,0 +1,232 @@
+"""
+SQL databases through their drivers: sqlite3, psycopg2 and psycopg (PostgreSQL), PyMySQL and mysql-connector
+(MySQL), asyncpg, SQLAlchemy's engines and text(), and Django's connection.cursor(). Each statement is read for what it does to which table, and a
+statement put together from the entry point's input is flagged.
+"""
+
+from __future__ import annotations
+
+import re
+
+from . import Call, Handle, Library, Use
+
+_CONNECT = {  # a driver's function that opens a connection -> the database, as people name it
+    "sqlite3.connect": "SQLite",
+    "psycopg2.connect": "PostgreSQL",
+    "psycopg.connect": "PostgreSQL",
+    "psycopg.Connection.connect": "PostgreSQL",
+    "psycopg.AsyncConnection.connect": "PostgreSQL",
+    "pymysql.connect": "MySQL",
+    "pymysql.connections.Connection": "MySQL",
+    "MySQLdb.connect": "MySQL",
+    "mysql.connector.connect": "MySQL",
+    "asyncpg.connect": "PostgreSQL",
+}
+
+_POOLS = {
+    "asyncpg.create_pool": "PostgreSQL",
+    "psycopg2.pool.SimpleConnectionPool": "PostgreSQL",
+    "psycopg2.pool.ThreadedConnectionPool": "PostgreSQL",
+    "psycopg_pool.ConnectionPool": "PostgreSQL",
+}
+
+_SCHEMES = (
+    ("postgres", "PostgreSQL"),
+    ("mysql", "MySQL"),
+    ("mariadb", "MySQL"),
+    ("sqlite", "SQLite"),
+    ("mssql", "SQL Server"),
+    ("oracle", "Oracle"),
+    ("snowflake", "Snowflake"),
+    ("redshift", "Redshift"),
+)
+
+_TIMEOUTS = ("connect_timeout", "timeout", "command_timeout", "connection_timeout", "read_timeout")
+
+_RUN = (
+    "execute",
+    "executemany",
+    "executescript",
+    "fetch",
+    "fetchrow",
+    "fetchval",
+    "exec_driver_sql",
+    "scalar",
+    "scalars",
+    "copy_records_to_table",
+    "copy_from",
+    "copy_expert",
+)
+
+_STATEMENT = re.compile(
+    r"^\s*(?:\(\s*)?(WITH|SELECT|INSERT|UPDATE|DELETE|MERGE|REPLACE|UPSERT|CREATE|DROP|ALTER|"
+    r"TRUNCATE|CALL|COPY|GRANT|VACUUM|PRAGMA|SHOW|EXPLAIN|BEGIN|COMMIT|ROLLBACK|SET)\b",
+    re.I,
+)
+
+_TABLE = {  # a statement -> where its table is named
+    "SELECT": r"\bFROM\s+([\w.\"`\[\]{}]+)",
+    "WITH": r"\bFROM\s+([\w.\"`\[\]{}]+)",
+    "INSERT": r"\bINTO\s+([\w.\"`\[\]{}]+)",
+    "REPLACE": r"\bINTO\s+([\w.\"`\[\]{}]+)",
+    "UPSERT": r"\bINTO\s+([\w.\"`\[\]{}]+)",
+    "MERGE": r"\bINTO\s+([\w.\"`\[\]{}]+)",
+    "UPDATE": r"^\s*UPDATE\s+(?:ONLY\s+)?([\w.\"`\[\]{}]+)",
+    "DELETE": r"\bFROM\s+([\w.\"`\[\]{}]+)",
+    "CREATE": r"\b(?:TABLE|INDEX|VIEW)\s+(?:IF\s+NOT\s+EXISTS\s+)?([\w.\"`\[\]{}]+)",
+    "DROP": r"\b(?:TABLE|INDEX|VIEW)\s+(?:IF\s+EXISTS\s+)?([\w.\"`\[\]{}]+)",
+    "ALTER": r"\bTABLE\s+([\w.\"`\[\]{}]+)",
+    "TRUNCATE": r"^\s*TRUNCATE\s+(?:TABLE\s+)?([\w.\"`\[\]{}]+)",
+    "COPY": r"^\s*COPY\s+([\w.\"`\[\]{}]+)",
+}
+
+_MARKERS = {"SQLite": "?", "PostgreSQL": "%s", "MySQL": "%s"}  # how a driver's parameters are written
+
+
+def _statement(text: str) -> tuple[str, str]:
+    """
+    What a SQL statement does and to which table, from its text (a literal or an f-string, as code):
+    "'SELECT id FROM orders WHERE ...'" -> ('SELECT', 'orders'). ('', '') when it isn't SQL it can read.
+    """
+    body = re.sub(r"^[rbuf]*(['\"]{1,3})", "", text.strip(), flags=re.I)
+    found = _STATEMENT.match(body)
+    if not found:
+        return "", ""
+    verb = found.group(1).upper()
+    if verb == "WITH":  # WITH ... AS (...) then the statement that does the work
+        main = re.search(r"\)\s*(SELECT|INSERT|UPDATE|DELETE|MERGE)\b", body, re.I)
+        verb = main.group(1).upper() if main else "SELECT"
+    table = re.search(_TABLE.get(verb, r"$^"), body, re.I | re.S)
+    return verb, table.group(1).strip('"`[]') if table else ""
+
+
+def _access(verb: str) -> str:
+    if verb in ("SELECT", "SHOW", "EXPLAIN", "PRAGMA"):
+        return "reads"
+    if verb == "CALL":
+        return "runs"
+    return "writes"
+
+
+class SqlLibrary(Library):
+    """
+    SQL through a DB-API driver (connect(), cursor(), execute()) or SQLAlchemy (create_engine(), connect(),
+    execute(text(...))).
+    """
+
+    name = "sql"
+    title = "SQL"
+    modules = (
+        "sqlite3",
+        "psycopg2",
+        "psycopg",
+        "psycopg_pool",
+        "pymysql",
+        "MySQLdb",
+        "mysql",
+        "asyncpg",
+        "sqlalchemy",
+        "django",
+    )
+    icon = "database"
+
+    def make(self, name: str, call: Call) -> Handle | Use | None:
+        if name in _CONNECT or name in _POOLS:
+            service = _CONNECT.get(name) or _POOLS[name]
+            target = call.shown(0, "database") if service == "SQLite" else ""
+            waits = (
+                ""
+                if service == "SQLite" or call.given(*_TIMEOUTS) or "timeout" in call.code(0, most=300)
+                else f"as long as {service} takes to answer"
+            )
+            kind = "connection" if name in _CONNECT else "pool"
+            facts = (("database", target),) if target else ()
+            fix = f"{name}(..., {'timeout' if name.startswith('asyncpg') else 'connect_timeout'}=5)"
+            return Handle(
+                kind,
+                service,
+                facts + (("driver", name.split(".")[0]),),
+                opens=service != "SQLite",
+                waits=waits,
+                fix=fix,
+            )  # SQLite opens a file: nothing to keep a connection to
+        if name == "django.db.connection.cursor":  # Django opens and keeps the connection itself
+            return Handle("cursor", "", (("driver", "django"),))
+        last = name.rsplit(".", 1)[-1]
+        if name.startswith("sqlalchemy") and last in ("create_engine", "create_async_engine", "engine_from_config"):
+            url = str(call.value(0, "url") or "")
+            service = next((label for scheme, label in _SCHEMES if url.startswith(scheme)), "SQL database")
+            return Handle("engine", service, (("driver", "sqlalchemy"),), opens=True)
+        if name.startswith("sqlalchemy") and last == "text":
+            return Handle(
+                "statement",
+                "",
+                (
+                    ("sql", call.text(0)),
+                    ("built", "yes" if call.built(0) else ""),
+                    ("inputs", ", ".join(call.inputs(0))),
+                ),
+            )
+        return None
+
+    def call(self, handle: Handle, method: str, call: Call) -> Handle | Use | None:
+        kind = handle.kind
+        if kind == "pool" and method in ("getconn", "acquire", "connection"):
+            return handle.but("connection")
+        if kind == "engine" and method in ("connect", "begin", "raw_connection"):
+            return handle.but("connection")
+        if kind == "connection" and method == "cursor":
+            return handle.but("cursor")
+        if kind == "connection" and method in ("transaction", "begin"):
+            return handle.but("connection")
+        if method not in _RUN:
+            return None
+        return self._statement(handle, method, call)
+
+    def _statement(self, handle: Handle, method: str, call: Call) -> Use:
+        service = handle.service or "SQL database"
+        held = call.held(0, "statement") or call.held(0, "query") or call.held(0, "sql")
+        if held is not None and held.kind == "statement":
+            text, built = held.fact("sql"), bool(held.fact("built"))
+            inputs = [part for part in held.fact("inputs").split(", ") if part]
+        else:
+            text = call.text(0, "sql") or call.text(0, "query") or call.text(0, "operation")
+            built = call.built(0, "sql") or call.built(0, "query") or call.built(0, "operation")
+            inputs = call.inputs(0, "sql") + call.inputs(0, "query") + call.inputs(0, "operation")
+        verb, table = _statement(text) if text else ("", "")
+        if method == "executescript":
+            verb, table = verb or "SCRIPT", table
+        operation = verb or "a query"
+        access = _access(verb) if verb else "runs"
+        flag = why = ""
+        if built and inputs:
+            marker = "$1" if handle.fact("driver") == "asyncpg" else _MARKERS.get(service, "%s")
+            flag = "SQL injection"
+            why = (
+                f"puts {', '.join(dict.fromkeys(inputs))} into the SQL text itself, so whoever sends it can change "
+                f"the statement (SQL injection). Pass it as a parameter instead: "
+                f'execute("... WHERE id = {marker}", (value,)).'
+            )
+        batch = ""
+        if method == "execute" and access == "writes" and verb in ("INSERT", "UPDATE", "DELETE", "REPLACE", "UPSERT"):
+            batch = "one execute() with a list of rows" if handle.fact("driver") == "sqlalchemy" else "executemany()"
+        elif access == "reads":
+            batch = "one query for all of them (WHERE ... IN (...), or a join)"
+        return Use(
+            operation,
+            access,
+            table,
+            service,
+            flag=flag,
+            why=why,
+            batch=batch,
+            gives=handle.but("cursor") if handle.kind == "connection" else None,
+        )
+
+    def describe(self, handle: Handle) -> str:
+        service = handle.service or "SQL"
+        return {
+            "statement": "a SQL statement",
+            "engine": f"a {service} engine",
+            "pool": f"a {service} connection pool",
+        }.get(handle.kind, f"a {service} {handle.kind}")