chore: import upstream snapshot with attribution

2026-07-13 12:32:25 +08:00
commit e014feafe1
2285 changed files with 1131979 additions and 0 deletions
@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+import os
+import sys
+import traceback
+import psycopg2
+
+# ---------------------------------------------------------------------------
+# Query list (kept at top for consistency with other tests)
+# ---------------------------------------------------------------------------
+
+TEST_QUERIES = [
+    "DROP TABLE IF EXISTS test",
+    "create table test (pk int, value int, d1 decimal(9, 3), f1 float, c1 char(10), t1 text, primary key(pk))",
+    "select * from test",
+    "insert into test (pk, value, d1, f1, c1, t1) values (0,0,0.0,0.0,'abc','a1')",
+    "select * from test",
+    "select dolt_add('-A');",
+    "select dolt_commit('-m', 'my commit')",
+    "select COUNT(*) FROM dolt.log",
+    "select dolt_checkout('-b', 'mybranch')",
+    "insert into test (pk, value, d1, f1, c1, t1) values (10,10, 123456.789, 420.42,'example','some text')",
+    "select dolt_commit('-a', '-m', 'my commit2')",
+    "select dolt_checkout('main')",
+    "select dolt_merge('mybranch')",
+    "select COUNT(*) FROM dolt.log",
+]
+
+# ---------------------------------------------------------------------------
+
+def env(name, default=None):
+    return os.getenv(name, default)
+
+
+def connect(user: str, port: int):
+    conn = psycopg2.connect(
+        host=env("PGHOST", "localhost"),
+        port=port,
+        dbname="postgres",
+        user=user,
+        password=env("PGPASSWORD", "password"),
+        connect_timeout=int(env("PGCONNECT_TIMEOUT", "10")),
+        sslmode=env("PGSSLMODE"),
+    )
+    conn.autocommit = True
+    return conn
+
+
+def run(cur, q):
+    print(f"SQL> {q}", flush=True)
+    cur.execute(q)
+    if cur.description is not None:
+        cur.fetchall()  # drain result set
+
+# load_test creates a table with |n_rows| and asserts that all rows are correctly returned.
+def load_test(cur, n_rows=1000):
+    print("\n=== Part 1: Load test ===", flush=True)
+
+    rows = max(1000, int(n_rows))
+
+    run(cur, "DROP TABLE IF EXISTS load_test")
+    run(cur, "CREATE TABLE load_test (id INT PRIMARY KEY, val INT NOT NULL)")
+
+    data = [(i, i * 10) for i in range(rows)]
+    cur.executemany(
+        "INSERT INTO load_test (id, val) VALUES (%s, %s)",
+        data,
+    )
+
+    cur.execute("SELECT COUNT(*) FROM load_test")
+    cnt = cur.fetchone()[0]
+    if cnt != rows:
+        raise AssertionError(f"COUNT(*) mismatch: expected {rows}, got {cnt}")
+
+    cur.execute("SELECT id FROM load_test ORDER BY id")
+    got = cur.fetchall()
+    if len(got) != rows:
+        raise AssertionError(f"fetchall mismatch: expected {rows}, got {len(got)}")
+
+    print(f"Inserted and selected {rows} rows OK.", flush=True)
+
+
+def compliance_test(cur):
+    print("\n=== Part 2: Test Queries ===", flush=True)
+    for q in TEST_QUERIES:
+        run(cur, q)
+    print("Compliance queries executed OK.", flush=True)
+
+
+def main():
+    if len(sys.argv) != 3:
+        print("Usage: python3 psycopg2_test.py <user> <port>")
+        return 2
+
+    user = sys.argv[1]
+    port = int(sys.argv[2])
+    load_rows = int(env("LOAD_ROWS", "1000"))
+
+    try:
+        with connect(user, port) as conn:
+            with conn.cursor() as cur:
+                load_test(cur, load_rows)
+                compliance_test(cur)
+
+        print("\n✅ All tests passed.", flush=True)
+        return 0
+
+    except Exception as e:
+        print("\n❌ Test failed.", flush=True)
+        print(f"Error: {e}", flush=True)
+        traceback.print_exc()
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())
@@ -0,0 +1,67 @@
+import sqlalchemy
+import psycopg2
+
+from sqlalchemy.engine import Engine
+from sqlalchemy import create_engine
+from sqlalchemy import text
+
+import sys
+
+QUERY_RESPONSE = [
+    {"create table test (pk int, \"value\" int, primary key(pk))": []},
+    {"describe test": [
+        ('pk', 'integer', 'NO', 'PRI', None, ''),
+        ('value', 'integer', 'YES', '', None, '')
+    ]},
+    {"insert into test (pk, value) values (0,0)": ()},
+    {"select * from test": [(0, 0)]},
+    {"select dolt_add('-A');": [(['0'],)]},
+    {"select dolt_commit('-m', 'my commit')": [('',)]},
+    {"select COUNT(*) FROM dolt_log": [(3,)]},
+    {"select dolt_checkout('-b', 'mybranch')": [(['0', "Switched to branch 'mybranch'"],)]},
+    {"insert into test (pk, value) values (1,1)": []},
+    {"select dolt_commit('-a', '-m', 'my commit2')": [('',)]},
+    {"select dolt_checkout('main')": [(['0', "Switched to branch 'main'"],)]},
+    {"select dolt_merge('mybranch')": [('',1,0,)]},
+    {"select COUNT(*) FROM dolt_log": [(4,)]},
+]
+
+
+def main():
+    user = sys.argv[1]
+    port = int(sys.argv[2])
+
+    conn_string_base = "postgresql+psycopg2://"
+
+    engine = create_engine(conn_string_base +
+                           "{user}:password@127.0.0.1:{port}/postgres".format(user=user,
+                                                                 port=port)
+                           )
+
+    with engine.connect() as con:
+        for query_response in QUERY_RESPONSE:
+            query = list(query_response.keys())[0]
+            exp_results = query_response[query]
+
+            result_proxy = con.execute(text(query))
+
+            try:
+                results = result_proxy.fetchall()
+                if (results != exp_results) and ("dolt_commit" not in query) and ("dolt_merge" not in query):
+                    print("Query:")
+                    print(query)
+                    print("Expected:")
+                    print(exp_results)
+                    print("Received:")
+                    print(results)
+                    sys.exit(1)
+            # You can't call fetchall on an insert
+            # so we'll just ignore the exception
+            except sqlalchemy.exc.ResourceClosedError:
+                pass
+
+    con.close()
+    sys.exit(0)
+
+
+main()