add foreign key support for users id and posts id using the "--foreign-keys" switch

WARNING: when using the foreign keys option, some entries in votes and postlinks might be updated to enforce data integrity

Author: madtibo

load_into_pg.py

@@ -29,24 +29,131 @@ def _createCmdTuple(cursor, keys, templ, attribs):
     """Use the cursor to mogrify a tuple of data.
     The passed data in `attribs` is augmented with default data (NULLs) and the
     order of data in the tuple is the same as in the list of `keys`. The
-    `cursor` is used toe mogrify the data and the `templ` is the template used
+    `cursor` is used to mogrify the data and the `templ` is the template used
     for the mogrification.
     """
     defs = _makeDefValues(keys)
     defs.update(attribs)
     return cursor.mogrify(templ, defs)
 
+def _getTableKeys(table):
+    """Return an array of the keys for a given table"""
+    keys = None
+    if table == 'Users':
+        keys = [
+            'Id'
+            , 'Reputation'
+            , 'CreationDate'
+            , 'DisplayName'
+            , 'LastAccessDate'
+            , 'WebsiteUrl'
+            , 'Location'
+            , 'AboutMe'
+            , 'Views'
+            , 'UpVotes'
+            , 'DownVotes'
+            , 'ProfileImageUrl'
+            , 'Age'
+            , 'AccountId'
+        ]
+    elif table == 'Badges':
+        keys = [
+            'Id'
+            , 'UserId'
+            , 'Name'
+            , 'Date'
+        ]
+    elif table == 'PostLinks':
+        keys = [
+            'Id'
+            , 'CreationDate'
+            , 'PostId'
+            , 'RelatedPostId'
+            , 'LinkTypeId'
+        ]
+    elif table == 'Comments':
+        keys = [
+            'Id'
+            , 'PostId'
+            , 'Score'
+            , 'Text'
+            , 'CreationDate'
+            , 'UserId'
+        ]
+    elif table == 'Votes':
+        keys = [
+            'Id'
+            , 'PostId'
+            , 'VoteTypeId'
+            , 'UserId'
+            , 'CreationDate'
+            , 'BountyAmount'
+        ]
+    elif table == 'Posts':
+        keys = [
+            'Id'
+            , 'PostTypeId'
+            , 'AcceptedAnswerId'
+            , 'ParentId'
+            , 'CreationDate'
+            , 'Score'
+            , 'ViewCount'
+            , 'Body'
+            , 'OwnerUserId'
+            , 'LastEditorUserId'
+            , 'LastEditorDisplayName'
+            , 'LastEditDate'
+            , 'LastActivityDate'
+            , 'Title'
+            , 'Tags'
+            , 'AnswerCount'
+            , 'CommentCount'
+            , 'FavoriteCount'
+            , 'ClosedDate'
+            , 'CommunityOwnedDate'
+        ]
+    elif table == 'Tags':
+        keys = [
+            'Id'
+            , 'TagName'
+            , 'Count'
+            , 'ExcerptPostId'
+            , 'WikiPostId'
+        ]
+    elif table == 'PostHistory':
+        keys = [
+            'Id',
+            'PostHistoryTypeId',
+            'PostId',
+            'RevisionGUID',
+            'CreationDate',
+            'UserId',
+            'Text'
+        ]
+    elif table == 'Comments':
+        keys = [
+            'Id',
+            'PostId',
+            'Score',
+            'Text',
+            'CreationDate',
+            'UserId',
+        ]
+    return keys
+
 def handleTable(table, keys, dbname, mbDbFile, mbHost, mbPort, mbUsername, mbPassword):
     """Handle the table including the post/pre processing."""
+    keys       = _getTableKeys(table)
     dbFile     = mbDbFile if mbDbFile is not None else table + '.xml'
     tmpl       = _createMogrificationTemplate(table, keys)
     start_time = time.time()
 
     try:
         pre    = open('./sql/' + table + '_pre.sql').read()
         post   = open('./sql/' + table + '_post.sql').read()
+        fk     = open('./sql/' + table + '_fk.sql').read()
     except IOError as e:
-        six.print_("Could not load pre/post sql. Are you running from the correct path?", file=sys.stderr)
+        six.print_("Could not load pre/post/fk sql. Are you running from the correct path?", file=sys.stderr)
         sys.exit(-1)
 
     dbConnectionParam = "dbname={}".format(dbname)
@@ -65,6 +172,7 @@ def handleTable(table, keys, dbname, mbDbFile, mbHost, mbPort, mbUsername, mbPas
     if mbPassword is not None:
         dbConnectionParam += ' password={}'.format(mbPassword)
 
+
     try:
         with pg.connect(dbConnectionParam) as conn:
             with conn.cursor() as cur:
@@ -86,13 +194,12 @@ def handleTable(table, keys, dbname, mbDbFile, mbHost, mbPort, mbUsername, mbPas
                                                 for row_attribs in rows
                                             ]
                                         )
-
                             if len(valuesStr) > 0:
                                 cmd = 'INSERT INTO ' + table + \
                                       ' VALUES\n' + valuesStr + ';'
                                 cur.execute(cmd)
                                 conn.commit()
-                        six.print_('Table processing took {:.1f} seconds'.format(time.time() - start_time))
+                        six.print_('Table processing took {1:.1f} seconds'.format(table, time.time() - start_time))
 
                         # Post-processing (creation of indexes)
                         start_time = time.time()
@@ -101,6 +208,14 @@ def handleTable(table, keys, dbname, mbDbFile, mbHost, mbPort, mbUsername, mbPas
                             cur.execute(post)
                             conn.commit()
                         six.print_('Post processing took {} seconds'.format(time.time() - start_time))
+                        if createFk:
+                            # fk-processing (creation of foreign keys)
+                            start_time = time.time()
+                            six.print_('fk processing ...')
+                            if post != '':
+                                cur.execute(fk)
+                                conn.commit()
+                            six.print_('fk processing took {} seconds'.format(time.time() - start_time))
 
                 except IOError as e:
                     six.print_("Could not read from file {}.".format(dbFile), file=sys.stderr)
@@ -113,8 +228,6 @@ def handleTable(table, keys, dbname, mbDbFile, mbHost, mbPort, mbUsername, mbPas
         six.print_("Warning from the database.", file=sys.stderr)
         six.print_("pg.Warning: {0}".format(str(w)), file=sys.stderr)
 
-
-
 #############################################################
 
 parser = argparse.ArgumentParser()
@@ -159,127 +272,30 @@ def handleTable(table, keys, dbname, mbDbFile, mbHost, mbPort, mbUsername, mbPas
                    , default = False
                    )
 
+parser.add_argument( '--foreign-keys'
+                   , help    = 'Create foreign keys.'
+                   , action = 'store_true'
+                   , default = False
+                   )
+
 args = parser.parse_args()
 
 table = args.table
-keys = None
-
-if table == 'Users':
-    keys = [
-        'Id'
-      , 'Reputation'
-      , 'CreationDate'
-      , 'DisplayName'
-      , 'LastAccessDate'
-      , 'WebsiteUrl'
-      , 'Location'
-      , 'AboutMe'
-      , 'Views'
-      , 'UpVotes'
-      , 'DownVotes'
-      , 'ProfileImageUrl'
-      , 'Age'
-      , 'AccountId'
-    ]
-elif table == 'Badges':
-    keys = [
-        'Id'
-      , 'UserId'
-      , 'Name'
-      , 'Date'
-    ]
-elif table == 'PostLinks':
-    keys = [
-        'Id'
-      , 'CreationDate'
-      , 'PostId'
-      , 'RelatedPostId'
-      , 'LinkTypeId'
-    ]
-elif table == 'Comments':
-    keys = [
-        'Id'
-      , 'PostId'
-      , 'Score'
-      , 'Text'
-      , 'CreationDate'
-      , 'UserId'
-    ]
-elif table == 'Votes':
-    keys = [
-        'Id'
-      , 'PostId'
-      , 'VoteTypeId'
-      , 'UserId'
-      , 'CreationDate'
-      , 'BountyAmount'
-    ]
-elif table == 'Posts':
-    keys = [
-        'Id'
-      , 'PostTypeId'
-      , 'AcceptedAnswerId'
-      , 'ParentId'
-      , 'CreationDate'
-      , 'Score'
-      , 'ViewCount'
-      , 'Body'
-      , 'OwnerUserId'
-      , 'LastEditorUserId'
-      , 'LastEditorDisplayName'
-      , 'LastEditDate'
-      , 'LastActivityDate'
-      , 'Title'
-      , 'Tags'
-      , 'AnswerCount'
-      , 'CommentCount'
-      , 'FavoriteCount'
-      , 'ClosedDate'
-      , 'CommunityOwnedDate'
-    ]
-
-    # If the user has not explicitly asked for loading the body, we replace it with NULL
-    if not args.with_post_body:
-        specialRules[('Posts', 'Body')] = 'NULL'
-
-elif table == 'Tags':
-    keys = [
-        'Id'
-      , 'TagName'
-      , 'Count'
-      , 'ExcerptPostId'
-      , 'WikiPostId'
-    ]
-elif table == 'PostHistory':
-    keys = [
-        'Id',
-        'PostHistoryTypeId',
-        'PostId',
-        'RevisionGUID',
-        'CreationDate',
-        'UserId',
-        'Text'
-    ]
-elif table == 'Comments':
-    keys = [
-        'Id',
-        'PostId',
-        'Score',
-        'Text',
-        'CreationDate',
-        'UserId',
-    ]
 
 try:
     # Python 2/3 compatibility
     input = raw_input
 except NameError:
     pass
 
-choice = input('This will drop the {} table. Are you sure [y/n]?'.format(table))
 
+if table == 'Posts':
+    # If the user has not explicitly asked for loading the body, we replace it with NULL
+    if not args.with_post_body:
+        specialRules[('Posts', 'Body')] = 'NULL'
+
+choice = input('This will drop the {} table. Are you sure [y/n]?'.format(table))
 if len(choice) > 0 and choice[0].lower() == 'y':
     handleTable(table, keys, args.dbname, args.file, args.host, args.port, args.username, args.password)
 else:
     six.print_("Cancelled.")
-

sql/Badges_fk.sql

@@ -0,0 +1 @@
+ALTER TABLE badges ADD CONSTRAINT fk_badges_userid FOREIGN KEY (userid) REFERENCES users (id);

sql/Comments_fk.sql

@@ -0,0 +1,2 @@
+ALTER TABLE Comments ADD CONSTRAINT fk_comments_userid FOREIGN KEY (userid) REFERENCES users (id);
+ALTER TABLE Comments ADD CONSTRAINT fk_comments_postid FOREIGN KEY (postid) REFERENCES posts (id);

sql/Comments_post.sql

@@ -6,4 +6,4 @@ CREATE INDEX cmnts_postid_idx ON Comments USING hash (PostId)
 CREATE INDEX cmnts_creation_date_idx ON Comments USING btree (CreationDate)
        WITH (FILLFACTOR = 100);
 CREATE INDEX cmnts_userid_idx ON Comments USING btree (UserId)
-       WITH (FILLFACTOR = 100);
+       WITH (FILLFACTOR = 100);

sql/PostHistory_fk.sql

@@ -0,0 +1,2 @@
+ALTER TABLE Posthistory ADD CONSTRAINT fk_posthistory_userid FOREIGN KEY (userid) REFERENCES users (id);
+ALTER TABLE Posthistory ADD CONSTRAINT fk_posthistory_postid FOREIGN KEY (postid) REFERENCES posts (id);

sql/PostLinks_fk.sql

@@ -0,0 +1,5 @@
+-- impossible to enforce so set NULL
+UPDATE Postlinks SET postid=NULL WHERE postid NOT IN (SELECT DISTINCT id FROM Posts);
+ALTER TABLE Postlinks ADD CONSTRAINT fk_postlinks_postid FOREIGN KEY (postid) REFERENCES posts (id);
+UPDATE Postlinks SET relatedpostid=NULL WHERE relatedpostid NOT IN (SELECT DISTINCT id FROM Posts);
+ALTER TABLE Postlinks ADD CONSTRAINT fk_postlinks_relatedpostid FOREIGN KEY (relatedpostid) REFERENCES posts (id);

sql/PostLinks_pre.sql

@@ -2,7 +2,7 @@ DROP TABLE IF EXISTS PostLinks CASCADE;
 CREATE TABLE PostLinks (
    Id                int         PRIMARY KEY ,
    CreationDate      timestamp   not NUll    ,
-   PostId            int         not NULL    ,
-   RelatedPostId     int         not NULL    ,
+   PostId            int         , -- not NULL    ,
+   RelatedPostId     int         , -- not NULL    ,
    LinkTypeId        int         not Null
 );

sql/Posts_fk.sql

@@ -0,0 +1,3 @@
+ALTER TABLE Posts ADD CONSTRAINT fk_posts_parentid FOREIGN KEY (parentid) REFERENCES posts (id);
+ALTER TABLE Posts ADD CONSTRAINT fk_posts_owneruserid FOREIGN KEY (owneruserid) REFERENCES users (id);
+ALTER TABLE Posts ADD CONSTRAINT fk_posts_lasteditoruserid FOREIGN KEY (lasteditoruserid) REFERENCES users (id);

sql/Tags_fk.sql

@@ -0,0 +1,2 @@
+-- dummy query
+SELECT 1;

sql/Tags_pre.sql

@@ -2,7 +2,7 @@ DROP TABLE IF EXISTS Tags CASCADE;
 CREATE TABLE Tags (
     Id                    int  PRIMARY KEY ,
     TagName               text not NULL    ,
-    Count                 int,
-    ExcerptPostId         int,
+    Count                 int              ,
+    ExcerptPostId         int              ,
     WikiPostId            int
 );

sql/Users_fk.sql

@@ -0,0 +1,2 @@
+-- dummy query
+SELECT 1;

sql/Votes_fk.sql

@@ -0,0 +1,4 @@
+ALTER TABLE Votes ADD CONSTRAINT fk_votes_userid FOREIGN KEY (userid) REFERENCES users (id);
+-- impossible to enforce so set NULL
+UPDATE Votes SET postid=NULL WHERE postid NOT IN (SELECT DISTINCT id FROM Posts);
+ALTER TABLE Votes ADD CONSTRAINT fk_votes_postid FOREIGN KEY (postid) REFERENCES posts (id);

sql/Votes_pre.sql

@@ -1,7 +1,7 @@
 DROP TABLE IF EXISTS Votes CASCADE;
 CREATE TABLE Votes (
    Id                int         PRIMARY KEY ,
-   PostId            int         not NULL    ,
+   PostId            int         , -- not NULL    ,
    VoteTypeId        int         not NULL    ,
    UserId            int                     ,
    CreationDate      timestamp   not NULL    ,