Source Code for Module osh.command.sql

  1  # osh 
  2  # Copyright (C) Jack Orenstein <jao@geophile.com> 
  3  # 
  4  # This program is free software; you can redistribute it and/or modify 
  5  # it under the terms of the GNU General Public License as published by 
  6  # the Free Software Foundation; either version 2 of the License, or 
  7  # (at your option) any later version. 
  8  # 
  9  # This program is distributed in the hope that it will be useful, 
 10  # but WITHOUT ANY WARRANTY; without even the implied warranty of 
 11  # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the 
 12  # GNU General Public License for more details. 
 13  # 
 14  # You should have received a copy of the GNU General Public License 
 15  # along with this program; if not, write to the Free Software 
 16  # Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA. 
 17   
 18  """C{sql [DB] QUERY} 
 19   
 20  Executes a sql query on a specified database.  Occurrences of 
 21  formatting directives (e.g. C{%s}) will be replaced by input values. 
 22   
 23  The database is selected by C{DB}. If C{DB} matches an C{osh.sql} 
 24  profile in C{.oshrc}, then the database is specified by that 
 25  profile. If C{DB} is omitted, then the default profile is used. 
 26   
 27  If C{QUERY} is a SELECT statement, then the query is executed and 
 28  output rows, represented by tuples, are written to output.  If C{QUERY} 
 29  is any other type of SQL statement, then no output is written. 
 30  """ 
 31   
 32  import os 
 33   
 34  import osh.loader 
 35  import osh.core 
 36   
 37  # CLI 
 38 -def _sql(): 
 39      return _Sql() 
 40   
 41  # CLI 
 42 -def sql(query, db = None): 
 43      """Executes a sql query on a specified database.  Occurrences of 
 44      formatting directives (e.g. C{%s}) will be replaced by input values. 
 45      The database is selected by C{db}. If C{db} matches an C{osh.sql} 
 46      profile in C{.oshrc}, then the database is specified by that 
 47      profile. If C{db} is C{None}, then the default profile is used. 
 48      If C{QUERY} is a I{select} statement, then the query is executed and 
 49      output rows, represented by tuples, are written to output. If C{QUERY} 
 50      is any other type of SQL statement, then no output is written. 
 51      """ 
 52      args = [] 
 53      if db: 
 54          args.append(db) 
 55      args.append(query) 
 56      return _Sql().process_args(*args) 
 57   
 58 -def _tuple_like(object): 
 59      return isinstance(object, list) or isinstance(object, tuple) 
 60   
 61 -class _Sql(osh.core.Generator): 
 62   
 63      _db_type = None 
 64      _host = None 
 65      _db = None 
 66      _header = None 
 67      _user = None 
 68      _password = None 
 69      _connection = None 
 70      _query = None 
 71   
 72   
 73      # object interface 
 74   
 75 -    def __init__(self): 
 76          osh.core.Generator.__init__(self, '', (1, 2)) 
 77   
 78   
 79      # OshCommand interface 
 80       
 81 -    def doc(self): 
 82          return __doc__ 
 83   
 84 -    def setup(self): 
 85          args = self.args() 
 86          db_profile = None 
 87          if args.has_next(): 
 88              query = args.next_string() 
 89              if args.has_next(): 
 90                  db_profile = query 
 91                  query = args.next_string() 
 92                  if args.has_next(): 
 93                      self.usage() 
 94          else: 
 95              self.usage() 
 96          if not db_profile: 
 97              db_profile_from_env = osh.core.default_db_profile 
 98              if db_profile_from_env: 
 99                  db_profile = db_profile_from_env 
100              else: 
101                  db_profile = osh.core.config_value('sql') 
102          if not db_profile: 
103              raise Exception('No db profile selected') 
104          # query 
105          self._query = query 
106          # schema 
107          if self.thread_state and type(self.thread_state) is osh.cluster.Host: 
108              schema = self.thread_state.schema 
109          else: 
110              schema = None 
111          # connection info 
112          driver = osh.core.config_value('sql', db_profile, 'driver') 
113          dbtype = osh.core.config_value('sql', db_profile, 'dbtype') 
114          if dbtype is None and driver is None: 
115              raise Exception('Database configuration osh.sql.%s must configure either dbtype or driver' % 
116                              db_profile) 
117          elif driver: 
118              connect_info = osh.core.config_subset('sql.%s' % db_profile) 
119              # Make sure set_schema specification is gone for establishing connection 
120              try: 
121                  del connect_info['set_schema'] 
122              except KeyError: 
123                  pass 
124              import sqladapter.dbapi 
125              self._db_type = sqladapter.dbapi.DBAPI() 
126              self._connection = self._db_type.connect(connect_info) 
127          elif dbtype: 
128              host = osh.core.config_value('sql', db_profile, 'host') 
129              if host is None: 
130                  host = 'localhost' 
131              db = osh.core.config_value('sql', db_profile, 'db') 
132              user = osh.core.config_value('sql', db_profile, 'user') 
133              password = osh.core.config_value('sql', db_profile, 'password') 
134              # Load dbtype module and connect 
135              self._db_type = osh.loader.load_and_create(dbtype) 
136              self._connection = self._db_type.connect(db, host, user, password) 
137          else: 
138              assert False 
139          set_schema_query = osh.core.config_value('sql', db_profile, 'set_schema') 
140          if set_schema_query and self.thread_state: 
141              schema = self.thread_state.schema 
142              self._db_type.run_update(self._connection, set_schema_query % schema, []) 
143   
144   
145      # Generator interface 
146   
147 -    def execute(self): 
148          # If this command is the first in a pipeline, then execute will be 
149          # called. 
150          self._execute_query([]) 
151   
152      # Receiver interface 
153   
154 -    def receive(self, object): 
155          self._execute_query(object) 
156   
157 -    def receive_complete(self): 
158          self._connection.close() 
159          self.send_complete() 
160   
161      # For use by this class 
162   
163 -    def _execute_query(self, inputs): 
164          for row in self._db_type.run_query(self._connection, 
165                                             self._query, 
166                                             inputs): 
167              self.send(row) 
168   
169 -class _DBType(object): 
170       
171 -    def connect(self, db, host, user, password): 
172          assert False 
173   
174 -    def run_query(self, connection, query, inputs): 
175          assert False 
176   
177 -    def close_connection(self, connection): 
178          assert False 
179