Wiring the Pieces Together

Length: 00:17:00

Lesson Summary:

With all of the functionality built into small, separate units, we're now ready to combine them to build our pgbackup tool. In this lesson, we'll create our main function that we can run when the tool is used.

Documentation for This Video

An Overview of the Parts

We've successfully written the following:

  • CLI parsing
  • Postgres Interaction
  • Local storage driver
  • AWS S3 storage driver

Now we want to create a function that we can use as a "console_script". Console scripts allow us to have setuptools create an executable script that we can install when someone installs our package.

Add "console_script" to project

We can make our project create a console script for us when a user runs pip install. This is similar to the way that we make executables before, except we don't need to manually do the work. To do this, we need to add an entry point in our setup.py:

setup.py (partial)

    package_dir={'': 'src'},
    install_requires=['boto3'],
    entry_points={
        'console_scripts': [
            'pgbackup=pgbackup.cli:main',
        ],
    }

Notice that we're referencing our cli module with a : and a main. That main is the function that we need to create now.

Wiring the Units Together

Our main function is going to go into the cli module, and it needs to do the following:

  1. Import the boto3 package.
  2. Import our pgdump and storage modules.
  3. Create a parser and parse the arguments.
  4. Fetch the database dump.
  5. Depending on the driver type, do either:
    • Create a boto3 S3 client and use storage.s3
      or
    • Open a local file and use storage.local

src/pgbackup/cli.py

def main():
    import boto3
    from pgbackup import pgdump, storage

    args = create_parser().parse_args()
    dump = pgdump.dump(args.url)
    if args.driver == 's3':
        client = boto3.client('s3')
        # TODO: create a better name based on the database name and the date
        storage.s3(client, dump.stdout, args.destination, 'example.sql')
    else:
        outfile = open(args.destination, 'wb')
        storage.local(dump.stdout, outfile)

Let's test it out:

$ pipenv shell
(pgbackup) $ pip install -e .
(pgbackup) $ pgbackup --driver local ./local-dump.sql postgres://demo:password@54.245.63.9:80/sample
(pgbackup) $ pgbackup --driver s3 pyscripting-db-backups postgres://demo:password@54.245.63.9:80/sample

Reviewing the Experience

It worked! That doesn't mean there aren't things to improve though. Here are some things we should fix:

  • Generate a good file name for S3
  • Create some output while the writing is happening
  • Create a shorthand switch for --driver (-d)

Generating a Dump File Name

For generating our file name, let's put all database URL interactions in the pgdump module with a function name of dump_file_name. We want the file name returned to be based on the database name, and it should also accept an optional timestamp. Let's work on the implementation now:

src/pgbackup/pgdump.py (partial)

def dump_file_name(url, timestamp=None):
    db_name = url.split("/")[-1]
    db_name = db_name.split("?")[0]
    if timestamp:
        return f"{db_name}-{timestamp}.sql"
    else:
        return f"{db_name}.sql"

Improving the CLI and Main Function

We want to add a shorthand -d flag to the driver argument, let's add that to the create_parser function:

src/pgbackup/cli.py (partial)

def create_parser():
    parser = argparse.ArgumentParser(description="""
    Back up PostgreSQL databases locally or to AWS S3.
    """)
    parser.add_argument("url", help="URL of database to backup")
    parser.add_argument("--driver", "-d",
            help="how & where to store backup",
            nargs=2,
            metavar=("DRIVER", "DESTINATION"),
            action=DriverAction,
            required=True)
    return parser

Lastly, let's generate a timestamp with time.strftime, generate a database file name, and print what we're doing as we upload/write files.

src/pgbackup/cli.py (partial)

def main():
    import time
    import boto3
    from pgbackup import pgdump, storage

    args = create_parser().parse_args()
    dump = pgdump.dump(args.url)

    if args.driver == 's3':
        client = boto3.client('s3')
        timestamp = time.strftime("%Y-%m-%dT%H:%M", time.localtime())
        file_name = pgdump.dump_file_name(args.url, timestamp)
        print(f"Backing database up to {args.destination} in S3 as {file_name}")
        storage.s3(client,
                dump.stdout,
                args.destination,
                file_name)
    else:
        outfile = open(args.destination, 'wb')
        print(f"Backing database up locally to {outfile.name}")
        storage.local(dump.stdout, outfile)

Feel free to test the CLI's modifications and commit these changes.


This lesson is only available to Linux Academy members.

Sign Up To View This Lesson
Or Log In

Looking For Team Training?

Learn More