Day 11 : psycopg2 操作

今天拉回 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()