今天拉回 python 來介紹 psycopg2，這個套件可以跟 postgres 進行互動。我們依賴該套件對 postgres 進行操作，就可以把 sqlstring 寫在 py 的腳本囉！最後我們也會介紹一下惡名昭彰的 SQL Injection 。

psycopg2 操作

1. Connection

import psycopg2

# Update connection string information
host = "localhost"
dbname = "XXX"
user = "XXX"
password = "XXX"
sslmode = "allow"

# Construct connection string
conn_string = "host={0} user={1} dbname={2} password={3} sslmode={4}".format(host, user, dbname, password, sslmode)
conn = psycopg2.connect(conn_string)
print("Connection established")

2. CREATE SCHEMA and INSERT

cursor = conn.cursor()

# Drop previous table of same name if one exists
cursor.execute("DROP TABLE IF EXISTS inventory;")
print("Finished dropping table (if existed)")

# Create a table
cursor.execute("CREATE TABLE inventory (id serial PRIMARY KEY, name VARCHAR(50), quantity INTEGER);")
print("Finished creating table")

# Insert some data into the table
cursor.execute("INSERT INTO inventory (name, quantity) VALUES (%s, %s);", ("banana", 150))
cursor.execute("INSERT INTO inventory (name, quantity) VALUES (%s, %s);", ("orange", 154))
cursor.execute("INSERT INTO inventory (name, quantity) VALUES (%s, %s);", ("apple", 100))
print("Inserted 3 rows of data")

# Clean up
conn.commit()
cursor.close()
conn.close()

小小提醒

什麼時候要 commit ?

當你需要進行資料庫修改的操作都需要，查詢沒有修改資料庫則不用

可以設定 conn.autocommit = True 就不用寫 commit 囉！

import psycopg2
conn = psycopg2.connect("dbname='db' user='user' host='localhost' password='password'")
conn.autocommit = True

或是

import psycopg2
conn = psycopg2.connect("dbname='db' user='user' host='localhost' password='password'")
con.set_session(autocommit=True)

3. Select data

# Construct connection string
conn_string = "host={0} user={1} dbname={2} password={3} sslmode={4}".format(host, user, dbname, password, sslmode)
conn = psycopg2.connect(conn_string)
print("Connection established")

cursor = conn.cursor()
# Fetch all rows from table
cursor.execute("SELECT * FROM inventory;")
rows = cursor.fetchall()

# Print all rows
for row in rows:
    print("Data row = (%s, %s, %s)" %(str(row[0]), str(row[1]), str(row[2])))

cursor.close()
conn.close()

小小提醒

• fetchall() 會一次取得 select 查詢的所有資料
  由於一次對資料庫進行全部資料讀取，情況適用於資料較小的時候；若資料庫資料過於龐大，會導致暫存記憶體不夠而無法讀取成功

  可以改使用 fetechmany 方法

• fetchone() 只會取一筆

• fetchmany(n) 取 n 筆

  while True:
    results = cursor.fetchmany(n)
    if not results:
        break
    for result in results:
        yield result
  

4. Update data

try:
    conn = psycopg2.connect(conn_string)
    cursor = conn.cursor()
    # Update a data row in the table
    cursor.execute("UPDATE inventory SET quantity = %s WHERE name = %s;", (200, "banana"))
    print("Updated 1 row of data")
except Exception as e:
    raise e
finally:
    cursor.close()
    conn.close()

5. Delete data

try:
    conn = psycopg2.connect(conn_string)
    cursor = conn.cursor()
    # Delete data row from table
    cursor.execute("DELETE FROM inventory WHERE name = %s;", ("orange",))
    print("Deleted 1 row of data")
except Exception as e:
    raise e
finally:
    cursor.close()
    conn.close()

6. Truncate Table

try:
    conn = psycopg2.connect(conn_string)
    cursor = conn.cursor()
    # Delete data row from table
    cursor.execute("TRUNCATE inventory;")
    print("TRUNCATE TABLE")
except Exception as e:
    raise e
finally:
    cursor.close()
    conn.close()

7. Drop Table

try:
    conn = psycopg2.connect(conn_string)
    cursor = conn.cursor()
    # Delete data row from table
    cursor.execute("DROP TABLE inventory;")
    print("DROP TABLE")
except Exception as e:
    raise e
finally:
    cursor.close()
    conn.close()

SQL Injection

SQL Injection 是 SQL 注入攻擊，是一種常見的攻擊侵害的手法

name = 'Erik'
result = cursor.execute(f""" select * from employee where name = '{name}' """)
print(result.fetchall())

若我們讓使用者自己設定名稱，也就是使用者可以動態去調整 SQL 敘述的結果，有心人士可以做...

故意讓 where 的條件永遠成立，就可以有意想不到的後果

name = "1' or '1'= '1"
result = cursor.execute(f""" select * from employee where name = '{name}' """)
print(result.fetchall())

正常的輸入方式

# Construct connection string
conn_string = "host={0} user={1} dbname={2} password={3} sslmode={4}".format(host, user, dbname, password, sslmode)
conn = psycopg2.connect(conn_string)
print("Connection established")

cursor = conn.cursor()
your_input = "apple"
# Fetch all rows from table
sqlstring = f'''
SELECT * FROM inventory where name = '{your_input}'
'''
print(sqlstring)
cursor.execute(sqlstring)
rows = cursor.fetchall()

# Print all rows
for row in rows:
    print("Data row = (%s, %s, %s)" %(str(row[0]), str(row[1]), str(row[2])))

cursor.close()
conn.close()

壞人的輸入方式

# Construct connection string
conn_string = "host={0} user={1} dbname={2} password={3} sslmode={4}".format(host, user, dbname, password, sslmode)
conn = psycopg2.connect(conn_string)
print("Connection established")

cursor = conn.cursor()
your_input = "1' or '1'= '1"
# Fetch all rows from table
sqlstring = f'''
SELECT * FROM inventory where name = '{your_input}'
'''
print(sqlstring)
cursor.execute(sqlstring)
rows = cursor.fetchall()

# Print all rows
for row in rows:
    print("Data row = (%s, %s, %s)" %(str(row[0]), str(row[1]), str(row[2])))

cursor.close()
conn.close()

更安全的輸入方式

# Construct connection string
conn_string = "host={0} user={1} dbname={2} password={3} sslmode={4}".format(host, user, dbname, password, sslmode)
conn = psycopg2.connect(conn_string)
print("Connection established")

cursor = conn.cursor()
your_input = "1' or '1'= '1"
# Fetch all rows from table
sqlstring = f'''
SELECT * FROM inventory where name = %s
'''
print(sqlstring, (your_input,))
cursor.execute(sqlstring, (your_input,))
rows = cursor.fetchall()

# Print all rows
for row in rows:
    print("Data row = (%s, %s, %s)" %(str(row[0]), str(row[1]), str(row[2])))

cursor.close()
conn.close()