Getting started with Python and psycopg2, I wanted to create a funtion that would insert a table row and return the primary key id column value.
The id would be returned whether the row existed already or not. This could easily be achieved with a PostgreSQL function, but this time I wanted to do it with Python.
Here is my sample table.<\/p>\n\n\n
\nDROP TABLE IF EXISTS crash.street;\n \nCREATE TABLE crash.street (\n id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,\n name TEXT NOT NULL,\n borough_code VARCHAR (2) NOT NULL,\n zip_code VARCHAR (5),\n create_date timestamp with time zone DEFAULT now(),\n update_date timestamp with time zone DEFAULT now(),\nCONSTRAINT s_nbz_u UNIQUE (name, borough_code, zip_code)\n);\n<\/pre><\/div>\n\n\nI created a Python module, ‘collision_db.py’, which contains all the functions for interacting with my database.<\/p>\n\n\n
\nimport psycopg2\nimport psycopg2.extras\nimport os, sys\n\n\ndef log_this(msg):\n print(f'{os.path.basename(__file__)}: {sys._getframe().f_back.f_lineno}> {msg}')\n\n\ndef get_db_connection():\n c = None\n my_schema = 'crash'\n\t# Set the search path so I don't have to specify \n\t# 'crash' schema in the SQL\n search_path = f'-c search_path=pg_catalog,public,{my_schema}'\n\n try:\n c = psycopg2.connect(\n database='nyc_data',\n user='postgres',\n password='xyz',\n options=search_path, )\n except psycopg2.DatabaseError as e:\n log_this(f'Error {e})')\n sys.exit(99)\n\n # No need for commits and rollbacks in for this test\n c.autocommit = True\n return c\n\n\ndef insert_into_street(*, con, street):\n street_id = None\n cur = con.cursor(cursor_factory=psycopg2.extras.DictCursor)\n cur.execute(\n "SELECT id FROM street WHERE name=%(name)s AND borough_code=%(borough_code)s AND zip_code=%(zip_code)s",\n street)\n log_this(f'Cursor status: {cur.statusmessage}')\n log_this(f'Cursor desc: {cur.description}')\n\n if cur.statusmessage == 'SELECT 0':\n log_this(f"Inserting new street {street}")\n cur.execute(\n "INSERT INTO street (name, borough_code, zip_code) VALUES ( %(name)s, %(borough_code)s, %(zip_code)s) RETURNING id",\n street)\n\n street_id = cur.fetchone()[0]\n return street_id\n\n...\n\ndef delete_from_table(*, con, table_name, id):\n cur = con.cursor()\n sql = f"DELETE FROM {table_name} WHERE id = %s"\n cur.execute(sql, (id,))\n return cur.rowcount\n\n<\/pre><\/div>\n\n\n
I will focus on the ‘insert_into_street’ function.<\/p>\n\n\n
\ndef insert_into_street(*, con, street):\n street_id = None\n cur = con.cursor(cursor_factory=psycopg2.extras.DictCursor)\n cur.execute(\n "SELECT id FROM street WHERE name=%(name)s AND borough_code=%(borough_code)s AND zip_code=%(zip_code)s",\n street)\n log_this(f'Cursor status: {cur.statusmessage}')\n log_this(f'Cursor desc: {cur.description}')\n\n if cur.statusmessage == 'SELECT 0':\n log_this(f"Inserting new street {street}")\n cur.execute(\n "INSERT INTO street (name, borough_code, zip_code) VALUES ( %(name)s, %(borough_code)s, %(zip_code)s) RETURNING id",\n street)\n\n street_id = cur.fetchone()[0]\n return street_id\n\n\n<\/pre><\/div>\n\n\nThe ‘insert_into_street’ is given a dictionary of street data, for example:<\/p>\n\n\n
\n expect_street = {\n 'name': 'Dead End Street',\n 'borough_code': 'q',\n 'zip_code': '11111'}\n<\/pre><\/div>\n\n\nThe function first checks if the street exists already. If it returns 0 rows, it inserts the the new street data using the PostgreSQL “RETURNING id”.
The cursor will now contain the ‘id’ value either from the select or the insert. This id will be returned by the function.<\/p>\n
To test this out I created a Pytest function that calls the insert function twice with the same data.
Both times it should return with the same row id. It should be noted that the ‘get_db_connection’ function
sets the Auto-commit to true.<\/p>\n\n\n\n
c.autocommit = True\n return c<\/code><\/pre>\n\n\n\nimport psycopg2\nimport pytest\n\nfrom collision_db import delete_from_table, get_db_connection, insert_into_street\nCON = None\n...\ndef setup_function():\n global CON\n CON = get_db_connection()\n log_this("In setup!")\n\ndef teardown_function():\n global CON\n CON.close()\n log_this("In teardown!")\n...\ndef test_insert_into_street_returns_street_id():\n global CON\n test_street = {\n 'name': 'dead_end',\n 'borough_code': 'm',\n 'zip_code': '11111'}\n street_id_1 = insert_into_street(con=CON, street=test_street)\n assert street_id_1 > 0\n street_id_2 = insert_into_street(con=CON, street=test_street)\n assert street_id_1 == street_id_2\n row_ct = delete_from_table(con=CON, table_name='street', id=street_id_2)\n assert row_ct == 1\n\n<\/pre><\/div>\n\n\nI run this test in noisy mode with -v and -s<\/p>\n\n\n
\n\u03bb pytest test_collision_db_1.py::test_insert_into_street_returns_street_id -v -s\n==================================================================== test session starts ==================================================================== platform win32 -- Python 3.7.2, pytest-5.4.2, py-1.8.1, pluggy-0.13.1 -- c:\\python37\\python.exe\ncachedir: .pytest_cache\nrootdir: C:\\Users\\ak1\\Apps\\Python\\collision\\tests\ncollected 1 item\n\ntest_collision_db_1.py::test_insert_into_street_returns_street_id test_collision_db_1.py: 25> In setup!\ncollision_db.py: 45> Cursor status: SELECT 0\ncollision_db.py: 46> Cursor desc: (Column(name='id', type_code=23),)\ncollision_db.py: 49> Inserting new street {'name': 'dead_end', 'borough_code': 'm', 'zip_code': '11111'}\ncollision_db.py: 45> Cursor status: SELECT 1\ncollision_db.py: 46> Cursor desc: (Column(name='id', type_code=23),)\nPASSED\ntest_collision_db_1.py: 30> In teardown!\n<\/pre><\/div>\n\n\nOr quiet mode.<\/p>\n\n\n
\n===================================================================== 1 passed in 0.19s =====================================================================\n\n#######################################\n\n#######################################\n\u03bb pytest test_collision_db_1.py::test_insert_into_street_returns_street_id \n==================================================================== test session starts ====================================================================\nplatform win32 -- Python 3.7.2, pytest-5.4.2, py-1.8.1, pluggy-0.13.1 \nrootdir: C:\\Users\\ak1\\Apps\\Python\\collision\\tests \ncollected 1 item \n \ntest_collision_db_1.py . [100%]\n \n===================================================================== 1 passed in 0.14s =====================================================================\n \n\n<\/pre><\/div>\n\n\nOutstanding!<\/p>\n","protected":false},"excerpt":{"rendered":"
Getting started with Python and psycopg2, I wanted to create a funtion that would insert a table row and return the primary key id column value. The id would be returned whether the row existed already or not. This could easily be achieved with a PostgreSQL function, but this time I wanted to do it […]<\/p>\n","protected":false},"author":1,"featured_media":0,"comment_status":"open","ping_status":"open","sticky":false,"template":"","format":"standard","meta":{"_jetpack_memberships_contains_paid_content":false,"footnotes":""},"categories":[97,77,100,98],"tags":[81,91,82,68,88],"class_list":["post-849","post","type-post","status-publish","format-standard","hentry","category-postgresql","category-psycopg2-python","category-python","category-sql","tag-postgresql","tag-psycopg2","tag-pytest","tag-python","tag-sql"],"jetpack_featured_media_url":"","jetpack_sharing_enabled":true,"_links":{"self":[{"href":"https:\/\/www.aibistin.com\/index.php?rest_route=\/wp\/v2\/posts\/849","targetHints":{"allow":["GET"]}}],"collection":[{"href":"https:\/\/www.aibistin.com\/index.php?rest_route=\/wp\/v2\/posts"}],"about":[{"href":"https:\/\/www.aibistin.com\/index.php?rest_route=\/wp\/v2\/types\/post"}],"author":[{"embeddable":true,"href":"https:\/\/www.aibistin.com\/index.php?rest_route=\/wp\/v2\/users\/1"}],"replies":[{"embeddable":true,"href":"https:\/\/www.aibistin.com\/index.php?rest_route=%2Fwp%2Fv2%2Fcomments&post=849"}],"version-history":[{"count":9,"href":"https:\/\/www.aibistin.com\/index.php?rest_route=\/wp\/v2\/posts\/849\/revisions"}],"predecessor-version":[{"id":859,"href":"https:\/\/www.aibistin.com\/index.php?rest_route=\/wp\/v2\/posts\/849\/revisions\/859"}],"wp:attachment":[{"href":"https:\/\/www.aibistin.com\/index.php?rest_route=%2Fwp%2Fv2%2Fmedia&parent=849"}],"wp:term":[{"taxonomy":"category","embeddable":true,"href":"https:\/\/www.aibistin.com\/index.php?rest_route=%2Fwp%2Fv2%2Fcategories&post=849"},{"taxonomy":"post_tag","embeddable":true,"href":"https:\/\/www.aibistin.com\/index.php?rest_route=%2Fwp%2Fv2%2Ftags&post=849"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}