jpayne@69: # Copyright 2002 by Andrew Dalke.  All rights reserved.
jpayne@69: # Revisions 2007-2010 copyright by Peter Cock.  All rights reserved.
jpayne@69: # Revisions 2009 copyright by Brad Chapman.  All rights reserved.
jpayne@69: # Revisions 2013 copyright by Tiago Antao.  All rights reserved.
jpayne@69: #
jpayne@69: # This file is part of the Biopython distribution and governed by your
jpayne@69: # choice of the "Biopython License Agreement" or the "BSD 3-Clause License".
jpayne@69: # Please see the LICENSE file that should have been included as part of this
jpayne@69: # package.
jpayne@69: #
jpayne@69: # Note that BioSQL (including the database schema and scripts) is
jpayne@69: # available and licensed separately.  Please consult www.biosql.org
jpayne@69: """Helper code for Biopython's BioSQL code (for internal use)."""
jpayne@69: 
jpayne@69: import os
jpayne@69: from typing import Dict, Type
jpayne@69: 
jpayne@69: 
jpayne@69: _dbutils: Dict[str, Type["Generic_dbutils"]] = {}
jpayne@69: 
jpayne@69: 
jpayne@69: class Generic_dbutils:
jpayne@69:     """Default database utilities."""
jpayne@69: 
jpayne@69:     def __init__(self):
jpayne@69:         """Create a Generic_dbutils object."""
jpayne@69: 
jpayne@69:     def tname(self, table):
jpayne@69:         """Return the name of the table."""
jpayne@69:         if table != "biosequence":
jpayne@69:             return table
jpayne@69:         else:
jpayne@69:             return "bioentry"
jpayne@69: 
jpayne@69:     def last_id(self, cursor, table):
jpayne@69:         """Return the last used id for a table."""
jpayne@69:         # XXX: Unsafe without transactions isolation
jpayne@69:         table = self.tname(table)
jpayne@69:         sql = f"select max({table}_id) from {table}"
jpayne@69:         cursor.execute(sql)
jpayne@69:         rv = cursor.fetchone()
jpayne@69:         return rv[0]
jpayne@69: 
jpayne@69:     def execute(self, cursor, sql, args=None):
jpayne@69:         """Just execute an sql command."""
jpayne@69:         cursor.execute(sql, args or ())
jpayne@69: 
jpayne@69:     def executemany(self, cursor, sql, seq):
jpayne@69:         """Execute many sql commands."""
jpayne@69:         cursor.executemany(sql, seq)
jpayne@69: 
jpayne@69:     def autocommit(self, conn, y=1):
jpayne@69:         """Set autocommit on the database connection."""
jpayne@69:         # Let's hope it was not really needed
jpayne@69: 
jpayne@69: 
jpayne@69: class Sqlite_dbutils(Generic_dbutils):
jpayne@69:     """Custom database utilities for SQLite."""
jpayne@69: 
jpayne@69:     def _sub_placeholder(self, sql):
jpayne@69:         """Format the argument placeholders for sqlite (PRIVATE)."""
jpayne@69:         return sql.replace("%s", "?")
jpayne@69: 
jpayne@69:     def execute(self, cursor, sql, args=None):
jpayne@69:         """Execute SQL command.
jpayne@69: 
jpayne@69:         Replaces %s with ? for variable substitution in sqlite3.
jpayne@69:         """
jpayne@69:         sql = self._sub_placeholder(sql)
jpayne@69:         cursor.execute(sql, args or ())
jpayne@69: 
jpayne@69:     def executemany(self, cursor, sql, seq):
jpayne@69:         """Execute many sql statements."""
jpayne@69:         sql = self._sub_placeholder(sql)
jpayne@69:         cursor.executemany(sql, seq)
jpayne@69: 
jpayne@69: 
jpayne@69: _dbutils["sqlite3"] = Sqlite_dbutils
jpayne@69: 
jpayne@69: 
jpayne@69: class Mysql_dbutils(Generic_dbutils):
jpayne@69:     """Custom database utilities for MySQL."""
jpayne@69: 
jpayne@69:     def last_id(self, cursor, table):
jpayne@69:         """Return the last used id for a table."""
jpayne@69:         if os.name == "java":
jpayne@69:             return Generic_dbutils.last_id(self, cursor, table)
jpayne@69:         try:
jpayne@69:             # This worked on older versions of MySQL
jpayne@69:             return cursor.insert_id()
jpayne@69:         except AttributeError:
jpayne@69:             # See bug 2390
jpayne@69:             # Google suggests this is the new way,
jpayne@69:             # same fix also suggested by Eric Gibert:
jpayne@69:             return cursor.lastrowid
jpayne@69: 
jpayne@69: 
jpayne@69: _dbutils["MySQLdb"] = Mysql_dbutils
jpayne@69: 
jpayne@69: 
jpayne@69: class _PostgreSQL_dbutils(Generic_dbutils):
jpayne@69:     """Base class for any PostgreSQL adaptor."""
jpayne@69: 
jpayne@69:     def next_id(self, cursor, table):
jpayne@69:         table = self.tname(table)
jpayne@69:         sql = f"SELECT nextval('{table}_pk_seq')"
jpayne@69:         cursor.execute(sql)
jpayne@69:         rv = cursor.fetchone()
jpayne@69:         return rv[0]
jpayne@69: 
jpayne@69:     def last_id(self, cursor, table):
jpayne@69:         table = self.tname(table)
jpayne@69:         sql = f"SELECT currval('{table}_pk_seq')"
jpayne@69:         cursor.execute(sql)
jpayne@69:         rv = cursor.fetchone()
jpayne@69:         return rv[0]
jpayne@69: 
jpayne@69: 
jpayne@69: class Psycopg2_dbutils(_PostgreSQL_dbutils):
jpayne@69:     """Custom database utilities for Psycopg2 (PostgreSQL)."""
jpayne@69: 
jpayne@69:     def autocommit(self, conn, y=True):
jpayne@69:         """Set autocommit on the database connection."""
jpayne@69:         if y:
jpayne@69:             if os.name == "java":
jpayne@69:                 conn.autocommit = 1
jpayne@69:             else:
jpayne@69:                 conn.set_isolation_level(0)
jpayne@69:         else:
jpayne@69:             if os.name == "java":
jpayne@69:                 conn.autocommit = 0
jpayne@69:             else:
jpayne@69:                 conn.set_isolation_level(1)
jpayne@69: 
jpayne@69: 
jpayne@69: _dbutils["psycopg2"] = Psycopg2_dbutils
jpayne@69: 
jpayne@69: 
jpayne@69: class Pgdb_dbutils(_PostgreSQL_dbutils):
jpayne@69:     """Custom database utilities for Pgdb (aka PyGreSQL, for PostgreSQL)."""
jpayne@69: 
jpayne@69:     def autocommit(self, conn, y=True):
jpayne@69:         """Set autocommit on the database connection. Currently not implemented."""
jpayne@69:         raise NotImplementedError("pgdb does not support this!")
jpayne@69: 
jpayne@69: 
jpayne@69: _dbutils["pgdb"] = Pgdb_dbutils
jpayne@69: 
jpayne@69: 
jpayne@69: def get_dbutils(module_name):
jpayne@69:     """Return the correct dbutils object for the database driver."""
jpayne@69:     try:
jpayne@69:         return _dbutils[module_name]()
jpayne@69:     except KeyError:
jpayne@69:         return Generic_dbutils()