Starter project files, async db chapter.

Author: mikeckennedy

code/ch8-async-databases/bin/load_data.py

@@ -0,0 +1,292 @@
+import json
+import os
+import sys
+import time
+from typing import List, Optional, Dict
+
+# noinspection PyPackageRequirements
+import progressbar
+from dateutil.parser import parse
+
+sys.path.insert(0, os.path.abspath(os.path.join(
+    os.path.dirname(__file__), "..")))
+
+import data.db_session as db_session
+from data.package import Package
+from data.release import Release
+from data.user import User
+
+
+def main():
+    init_db()
+    session = db_session.create_session()
+    user_count = session.query(User).count()
+    session.close()
+    if user_count == 0:
+        file_data = do_load_files()
+        users = find_users(file_data)
+
+        db_users = do_user_import(users)
+        do_import_packages(file_data, db_users)
+
+    do_summary()
+
+
+def do_summary():
+    session = db_session.create_session()
+
+    print("Final numbers:")
+    print("Users: {:,}".format(session.query(User).count()))
+    print("Packages: {:,}".format(session.query(Package).count()))
+    print("Releases: {:,}".format(session.query(Release).count()))
+
+
+def do_user_import(user_lookup: Dict[str, str]) -> Dict[str, User]:
+    print("Importing users ... ", flush=True)
+    with progressbar.ProgressBar(max_value=len(user_lookup)) as bar:
+        for idx, (email, name) in enumerate(user_lookup.items()):
+            session = db_session.create_session()
+            session.expire_on_commit = False
+
+            user = User()
+            user.email = email
+            user.name = name
+            session.add(user)
+
+            session.commit()
+            bar.update(idx)
+
+    print()
+    sys.stderr.flush()
+    sys.stdout.flush()
+
+    session = db_session.create_session()
+    return {u.email: u for u in session.query(User)}
+
+
+def do_import_packages(file_data: List[dict], user_lookup: Dict[str, User]):
+    errored_packages = []
+    print("Importing packages and releases ... ", flush=True)
+    with progressbar.ProgressBar(max_value=len(file_data)) as bar:
+        for idx, p in enumerate(file_data):
+            try:
+                load_package(p, user_lookup)
+                bar.update(idx)
+            except Exception as x:
+                errored_packages.append((p, " *** Errored out for package {}, {}".format(p.get('package_name'), x)))
+                raise
+    sys.stderr.flush()
+    sys.stdout.flush()
+    print()
+    print("Completed packages with {} errors.".format(len(errored_packages)))
+    for (p, txt) in errored_packages:
+        print(txt)
+
+
+def do_load_files() -> List[dict]:
+    data_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../data/pypi-top-100'))
+    print("Loading files from {}".format(data_path))
+    files = get_file_names(data_path)
+    print("Found {:,} files, loading ...".format(len(files)), flush=True)
+    time.sleep(.1)
+
+    file_data = []
+    with progressbar.ProgressBar(max_value=len(files)) as bar:
+        for idx, f in enumerate(files):
+            file_data.append(load_file_data(f))
+            bar.update(idx)
+
+    sys.stderr.flush()
+    sys.stdout.flush()
+    print()
+    return file_data
+
+
+def find_users(data: List[dict]) -> dict:
+    print("Discovering users...", flush=True)
+    found_users = {}
+
+    with progressbar.ProgressBar(max_value=len(data)) as bar:
+        for idx, p in enumerate(data):
+            info = p.get('info')
+            found_users.update(get_email_and_name_from_text(info.get('author'), info.get('author_email')))
+            found_users.update(get_email_and_name_from_text(info.get('maintainer'), info.get('maintainer_email')))
+            bar.update(idx)
+
+    sys.stderr.flush()
+    sys.stdout.flush()
+    print()
+    print("Discovered {:,} users".format(len(found_users)))
+    print()
+
+    return found_users
+
+
+def get_email_and_name_from_text(name: str, email: str) -> dict:
+    data = {}
+
+    if not name or not email:
+        return data
+
+    emails = email.strip().lower().split(',')
+    names = name
+    if len(email) > 1:
+        names = name.strip().split(',')
+
+    for n, e in zip(names, emails):
+        if not n or not e:
+            continue
+
+        data[e.strip()] = n.strip()
+
+    return data
+
+
+def load_file_data(filename: str) -> dict:
+    try:
+        with open(filename, 'r', encoding='utf-8') as fin:
+            data = json.load(fin)
+    except Exception as x:
+        print("ERROR in file: {}, details: {}".format(filename, x), flush=True)
+        raise
+
+    return data
+
+
+def load_package(data: dict, user_lookup: Dict[str, User]):
+    try:
+        info = data.get('info', {})
+
+        p = Package()
+        p.id = data.get('package_name', '').strip()
+        if not p.id:
+            return
+
+        p.author = info.get('author')
+        p.author_email = info.get('author_email')
+
+        releases = build_releases(p.id, data.get("releases", {}))
+
+        if releases:
+            p.created_date = releases[0].created_date
+
+        maintainers_lookup = get_email_and_name_from_text(info.get('maintainer'), info.get('maintainer_email'))
+        maintainers = []
+
+        p.summary = info.get('summary')
+        p.description = info.get('description')
+
+        p.home_page = info.get('home_page')
+        p.docs_url = info.get('docs_url')
+        p.package_url = info.get('package_url')
+
+        p.author = info.get('author')
+        p.author_email = info.get('author_email')
+        p.license = detect_license(info.get('license'))
+
+        session = db_session.create_session()
+        session.add(p)
+        session.add_all(releases)
+        if maintainers:
+            session.add_all(maintainers)
+        session.commit()
+        session.close()
+    except OverflowError:
+        # What the heck, people just putting fake data in here
+        # Size is terabytes...
+        pass
+    except Exception:
+        raise
+
+
+def detect_license(license_text: str) -> Optional[str]:
+    if not license_text:
+        return None
+
+    license_text = license_text.strip()
+
+    if len(license_text) > 100 or '\n' in license_text:
+        return "CUSTOM"
+
+    license_text = license_text \
+        .replace('Software License', '') \
+        .replace('License', '')
+
+    if '::' in license_text:
+        # E.g. 'License :: OSI Approved :: Apache Software License'
+        return license_text \
+            .split(':')[-1] \
+            .replace('  ', ' ') \
+            .strip()
+
+    return license_text.strip()
+
+
+def build_releases(package_id: str, releases: dict) -> List[Release]:
+    db_releases = []
+    for k in releases.keys():
+        all_releases_for_version = releases.get(k)
+        if not all_releases_for_version:
+            continue
+
+        v = all_releases_for_version[-1]
+
+        r = Release()
+        r.package_id = package_id
+        r.major_ver, r.minor_ver, r.build_ver = make_version_num(k)
+        r.created_date = parse(v.get('upload_time'))
+        r.comment = v.get('comment_text')
+        r.url = v.get('url')
+        r.size = int(v.get('size', 0))
+
+        db_releases.append(r)
+
+    return db_releases
+
+
+def make_version_num(version_text):
+    major, minor, build = 0, 0, 0
+    if version_text:
+        version_text = version_text.split('b')[0]
+        parts = version_text.split('.')
+        if len(parts) == 1:
+            major = try_int(parts[0])
+        elif len(parts) == 2:
+            major = try_int(parts[0])
+            minor = try_int(parts[1])
+        elif len(parts) == 3:
+            major = try_int(parts[0])
+            minor = try_int(parts[1])
+            build = try_int(parts[2])
+
+        return major, minor, build
+
+
+def try_int(text) -> int:
+    try:
+        return int(text)
+    except:
+        return 0
+
+
+def init_db():
+    top_folder = os.path.dirname(__file__)
+    rel_file = os.path.join('..', 'db', 'pypi.sqlite')
+    db_file = os.path.abspath(os.path.join(top_folder, rel_file))
+    db_session.global_init(db_file)
+
+
+def get_file_names(data_path: str) -> List[str]:
+    files = []
+    for f in os.listdir(data_path):
+        if f.endswith('.json'):
+            files.append(
+                os.path.abspath(os.path.join(data_path, f))
+            )
+
+    files.sort()
+    return files
+
+
+if __name__ == '__main__':
+    main()

code/ch8-async-databases/data/__all_models.py

@@ -0,0 +1,6 @@
+# noinspection PyUnresolvedReferences
+from data.package import Package
+# noinspection PyUnresolvedReferences
+from data.user import User
+# noinspection PyUnresolvedReferences
+from data.release import Release

code/ch8-async-databases/data/db_session.py

@@ -0,0 +1,50 @@
+from pathlib import Path
+from typing import Callable, Optional
+
+import sqlalchemy as sa
+import sqlalchemy.orm as orm
+from sqlalchemy.orm import Session
+
+from data.modelbase import SqlAlchemyBase
+
+__factory: Optional[Callable[[], Session]] = None
+
+
+def global_init(db_file: str):
+    global __factory
+
+    if __factory:
+        return
+
+    if not db_file or not db_file.strip():
+        raise Exception("You must specify a db file.")
+
+    folder = Path(db_file).parent
+    folder.mkdir(parents=True, exist_ok=True)
+
+    conn_str = 'sqlite:///' + db_file.strip()
+    print("Connecting to DB with {}".format(conn_str))
+
+    # Adding check_same_thread = False after the recording. This can be an issue about
+    # creating / owner thread when cleaning up sessions, etc. This is a sqlite restriction
+    # that we probably don't care about in this example.
+    engine = sa.create_engine(conn_str, echo=False, connect_args={"check_same_thread": False})
+    __factory = orm.sessionmaker(bind=engine)
+
+    # noinspection PyUnresolvedReferences
+    import data.__all_models
+
+    SqlAlchemyBase.metadata.create_all(engine)
+
+
+def create_session() -> Session:
+    global __factory
+
+    if not __factory:
+        raise Exception("You must call global_init() before using this method.")
+
+    session: Session = __factory()
+    session.expire_on_commit = False
+
+    return session
+

code/ch8-async-databases/data/modelbase.py

@@ -0,0 +1,3 @@
+import sqlalchemy.ext.declarative
+
+SqlAlchemyBase = sqlalchemy.ext.declarative.declarative_base()

code/ch8-async-databases/data/package.py

@@ -0,0 +1,44 @@
+import datetime
+from typing import List
+
+import sqlalchemy as sa
+import sqlalchemy.orm as orm
+from data.modelbase import SqlAlchemyBase
+from data.release import Release
+
+
+class Package(SqlAlchemyBase):
+    __tablename__ = 'packages'
+
+    id: str = sa.Column(sa.String, primary_key=True)
+    created_date: datetime.datetime = sa.Column(sa.DateTime, default=datetime.datetime.now, index=True)
+    last_updated: datetime.datetime = sa.Column(sa.DateTime, default=datetime.datetime.now, index=True)
+    summary: str = sa.Column(sa.String, nullable=False)
+    description: str = sa.Column(sa.String, nullable=True)
+
+    home_page: str = sa.Column(sa.String)
+    docs_url: str = sa.Column(sa.String)
+    package_url: str = sa.Column(sa.String)
+
+    author_name: str = sa.Column(sa.String)
+    author_email: str = sa.Column(sa.String, index=True)
+
+    license: str = sa.Column(sa.String, index=True)
+
+    # releases relationship
+    releases: List[Release] = orm.relation("Release", order_by=[
+        Release.major_ver.desc(),
+        Release.minor_ver.desc(),
+        Release.build_ver.desc(),
+    ], back_populates='package')
+
+    def __repr__(self):
+        return '<Package {}>'.format(self.id)
+
+
+# p = Package()  # one query
+#
+# print(p.id)
+# print("All releases")
+# for r in p.releases:
+#     print("{}.{}.{}".format(r.major_ver, r.minor_ver, r.build_ver))