1 files changed, 16 insertions, 6 deletions
diff --git a/bin/common/srtool_common.py b/bin/common/srtool_common.py
index d9fbd341..2a92333a 100755
--- a/bin/common/srtool_common.py
+++ b/bin/common/srtool_common.py
@@ -40,8 +40,10 @@ from datetime import datetime
 # Load the srt.sqlite schema index file
 # Since it is generated from this script
 # it may not exist on the first pass
+dir_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
+sys.path.insert(0, dir_path)
 try:
-    from srt_schema import ORM
+    from common.srt_schema import ORM
 except:
     # Do a pass so that '--generate-schema-header' can fix it
     print("Warning: srt_schema not yet created or bad format")
@@ -857,24 +859,29 @@ def update_cve_status_tree(cve_list,update_skip_history):
 # CREATE TABLE "orm_notifycategories" ("id" integer NOT NULL PRIMARY KEY AUTOINCREMENT, "category" varchar(50) NULL);
 # ...
 
-def gen_schema_header():
+
+def gen_schema_header(database_dir,schema_dir):
+
+    database_file = os.path.join(database_dir, 'srt.sqlite')
+    schema_file = os.path.join(schema_dir, 'srt_schema.py')
+
     create_re = re.compile(r"CREATE TABLE[A-Z ]* \"(\w+)\" \((.+)\);")
     try:
-        cmd = ('sqlite3', os.path.join(srtool_basepath, 'srt.sqlite'), '.schema')
+        cmd = ('sqlite3', database_file, '.schema')
         output = subprocess.check_output(cmd, stderr=subprocess.STDOUT)
     except subprocess.CalledProcessError as e:
         print("ERROR(%d): %s" % (e.returncode, e.output))
         return
 
     # Fetch USER_SRTOOL_ID
-    conn = sqlite3.connect(srtDbName)
+    conn = sqlite3.connect(database_file)
     cur = conn.cursor()
     USER_SRTOOL_NAME = 'SRTool'
     user = cur.execute("SELECT * FROM users_srtuser where username = '%s'" % USER_SRTOOL_NAME).fetchone()
     USER_SRTOOL_ID = user[0]    # Hardcoded 'ORM.USERS_SRTUSER_ID'
     conn.close()
 
-    with open(os.path.join(srtool_basepath,'bin/common/srt_schema.py'), 'w') as fd:
+    with open(schema_file, 'w') as fd:
         fd.write("# SRTool database table schema indexes\n")
         fd.write("# Generated by: './bin/common/srtool_common.py --generate-schema-header'\n")
         fd.write("# Should be run after any schema changes to sync commandline tools\n")
@@ -1089,6 +1096,7 @@ def main(argv):
     parser.add_argument('--init-notify-categories', '-n', action='store_const', const='init_notify_categories', dest='command', help='Initialize notify categories')
     parser.add_argument('--score-new-cves', '-s', dest='score_new_cves', help='Score CVEs for triage [NEW|CVE-1234]')
     parser.add_argument('--generate-schema-header', '-g', action='store_const', const='gen_schema_header', dest='command', help='Generate database schema header')
+    parser.add_argument('--generate-schema-header-dir', dest='gen_schema_header_dir', help='Generate database schema header for a give database directory')
 
 
     parser.add_argument('--update-cve-status-tree', '-S', dest='update_cve_status_tree', help="Update CVEs and their children's cumulative status")
@@ -1126,7 +1134,9 @@ def main(argv):
     elif args.score_new_cves:
         score_new_cves(args.score_new_cves)
     elif 'gen_schema_header' == args.command:
-        gen_schema_header()
+        gen_schema_header(srtool_basepath,os.path.join(srtool_basepath,'bin/common'))
+    elif args.gen_schema_header_dir:
+        gen_schema_header(args.gen_schema_header_dir,args.gen_schema_header_dir)
 
     elif args.update_cve_status_tree:
         update_cve_status_tree(args.update_cve_status_tree, update_skip_history)